Use discriminated unions where possible

This allows Pydantic to generate more useful validation errors.
dandi · Jun 21, 2024 · 6edc06a · 6edc06a
1 parent e135307
commit 6edc06a
Showing 1 changed file with 24 additions and 6 deletions.
diff --git a/dandischema/models.py b/dandischema/models.py
@@ -22,6 +22,9 @@
 )
 from pydantic.json_schema import JsonSchemaValue
 from pydantic_core import CoreSchema
+from typing_extensions import (
+    Annotated,  # TODO: import from `typing` when Python 3.8 support is dropped
+)
 from zarr_checksum.checksum import InvalidZarrChecksum, ZarrDirectoryDigest
 
 from .consts import DANDI_SCHEMA_VERSION
@@ -1188,9 +1191,14 @@ class Activity(DandiBaseModel):
 
     # isPartOf: Optional["Activity"] = Field(None, json_schema_extra={"nskey": "schema"})
     # hasPart: Optional["Activity"] = Field(None, json_schema_extra={"nskey": "schema"})
-    wasAssociatedWith: Optional[List[Union[Person, Organization, Software, Agent]]] = (
-        Field(None, json_schema_extra={"nskey": "prov"})
-    )
+    wasAssociatedWith: Optional[
+        List[
+            Annotated[
+                Union[Person, Organization, Software, Agent],
+                Field(discriminator="schemaKey"),
+            ]
+        ]
+    ] = Field(None, json_schema_extra={"nskey": "prov"})
     used: Optional[List[Equipment]] = Field(
         None,
         description="A listing of equipment used for the activity.",
@@ -1467,13 +1475,21 @@ class CommonModel(DandiBaseModel):
         description="A description of the item.",
         json_schema_extra={"nskey": "schema"},
     )
-    contributor: Optional[List[Union[Person, Organization]]] = Field(
+    contributor: Optional[
+        List[Annotated[Union[Person, Organization], Field(discriminator="schemaKey")]]
+    ] = Field(
         None,
         title="Contributors",
         description="Contributors to this item: persons or organizations.",
         json_schema_extra={"nskey": "schema"},
     )
-    about: Optional[List[Union[Disorder, Anatomy, GenericType]]] = Field(
+    about: Optional[
+        List[
+            Annotated[
+                Union[Disorder, Anatomy, GenericType], Field(discriminator="schemaKey")
+            ]
+        ]
+    ] = Field(
         None,
         title="Subject matter of the dataset",
         description="The subject matter of the content, such as disorders, brain anatomy.",
@@ -1586,7 +1602,9 @@ def contributor_musthave_contact(
         max_length=3000,
         json_schema_extra={"nskey": "schema"},
     )
-    contributor: List[Union[Person, Organization]] = Field(
+    contributor: List[
+        Annotated[Union[Person, Organization], Field(discriminator="schemaKey")]
+    ] = Field(
         title="Dandiset contributors",
         description="People or Organizations that have contributed to this Dandiset.",
         json_schema_extra={"nskey": "schema"},