Merge pull request #245 from dandi/pydantic-discriminated-unions

Use discriminated unions to provide more helpful error messages
dandi · Jul 10, 2024 · 06c0ee3 · 06c0ee3
2 parents e135307 + abccf00
commit 06c0ee3
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 8 deletions.
diff --git a/dandischema/consts.py b/dandischema/consts.py
@@ -1,4 +1,4 @@
-DANDI_SCHEMA_VERSION = "0.6.7"
+DANDI_SCHEMA_VERSION = "0.6.8"
 ALLOWED_INPUT_SCHEMAS = [
     "0.4.4",
     "0.5.1",
@@ -10,6 +10,7 @@
     "0.6.4",
     "0.6.5",
     "0.6.6",
+    "0.6.7",
 ]
 
 # ATM we allow only for a single target version which is current

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"},

diff --git a/dandischema/tests/data/metadata/meta_000008.json b/dandischema/tests/data/metadata/meta_000008.json
@@ -8,6 +8,7 @@
     {
       "schemaKey": "Person",
       "name": "Scala, Federico",
+      "email": "fscala@example.com",
       "roleName": [
         "dcite:DataCollector",
         "dcite:Author",
@@ -150,6 +151,7 @@
       "schemaKey": "Person",
       "identifier": "0000-0002-4305-6376",
       "name": "Tolias, Andreas Savas",
+      "email": "atolias@example.com",
       "roleName": [
         "dcite:Author",
         "dcite:ContactPerson"

diff --git a/dandischema/tests/test_datacite.py b/dandischema/tests/test_datacite.py
@@ -72,6 +72,7 @@ def metadata_basic() -> Dict[str, Any]:
                 "name": "A_last, A_first",
                 "email": "nemo@example.com",
                 "roleName": [RoleType("dcite:ContactPerson")],
+                "schemaKey": "Person",
             }
         ],
         "license": [LicenseType("spdx:CC-BY-4.0")],
@@ -195,12 +196,18 @@ def test_datacite(dandi_id: str, schema: Any) -> None:
                     {
                         "name": "A_last, A_first",
                         "roleName": [RoleType("dcite:ContactPerson")],
+                        "email": "nemo@example.com",
+                        "schemaKey": "Person",
                     },
                     {
                         "name": "B_last, B_first",
                         "roleName": [RoleType("dcite:Author")],
+                        "schemaKey": "Person",
+                    },
+                    {
+                        "name": "C_last, C_first",
+                        "schemaKey": "Person",
                     },
-                    {"name": "C_last, C_first"},
                 ],
             },
             {
@@ -218,10 +225,13 @@ def test_datacite(dandi_id: str, schema: Any) -> None:
                     {
                         "name": "A_last, A_first",
                         "roleName": [RoleType("dcite:ContactPerson")],
+                        "email": "nemo@example.com",
+                        "schemaKey": "Person",
                     },
                     {
                         "name": "B_last, B_first",
                         "roleName": [RoleType("dcite:Sponsor")],
+                        "schemaKey": "Person",
                     },
                 ],
             },
@@ -237,11 +247,14 @@ def test_datacite(dandi_id: str, schema: Any) -> None:
                     {
                         "name": "A_last, A_first",
                         "roleName": [RoleType("dcite:ContactPerson")],
+                        "email": "nemo@example.com",
+                        "schemaKey": "Person",
                     },
                     {
                         "name": "B_last, B_first",
                         "identifier": "0000-0001-0000-0000",
                         "roleName": [RoleType("dcite:Sponsor")],
+                        "schemaKey": "Person",
                     },
                 ],
             },
@@ -264,11 +277,14 @@ def test_datacite(dandi_id: str, schema: Any) -> None:
                     {
                         "name": "A_last, A_first",
                         "roleName": [RoleType("dcite:ContactPerson")],
+                        "email": "nemo@example.com",
+                        "schemaKey": "Person",
                     },
                     {
                         "name": "B_last, B_first",
                         "identifier": "0000-0001-0000-0000",
                         "roleName": [RoleType("dcite:Funder")],
+                        "schemaKey": "Person",
                     },
                 ],
             },
@@ -297,10 +313,13 @@ def test_datacite(dandi_id: str, schema: Any) -> None:
                             RoleType("dcite:Software"),
                         ],
                         "identifier": "0000-0001-0000-0000",
+                        "schemaKey": "Person",
                     },
                     {
                         "name": "B_last, B_first",
                         "roleName": [RoleType("dcite:ContactPerson")],
+                        "email": "nemo@example.com",
+                        "schemaKey": "Person",
                     },
                 ],
             },

diff --git a/dandischema/tests/test_models.py b/dandischema/tests/test_models.py
@@ -381,6 +381,7 @@ def test_dantimeta_1() -> None:
                 "name": "last name, first name",
                 "email": "someone@dandiarchive.org",
                 "roleName": [RoleType("dcite:ContactPerson")],
+                "schemaKey": "Person",
             }
         ],
         "license": [LicenseType("spdx:CC-BY-4.0")],