fix: DataDescriptionUpgrade sets data_level to raw if not correct (#568)

* fix: DataDescriptionUpgrade sets data_level to raw if field value is not correct

* test: add test with missing data_level field

* validator for data level

* fix: updates data level validator

* fix: adds check if user sets data_level explicitly

---------

Co-authored-by: Mekhla Kapoor <[email protected]>
Co-authored-by: jtyoung84 <[email protected]>

Author: 3 people

src/aind_data_schema/data_description.py

@@ -5,9 +5,9 @@
 import re
 from datetime import datetime
 from enum import Enum, EnumMeta
-from typing import Any, List, Optional
+from typing import Any, List, Optional, Union
 
-from pydantic import Field
+from pydantic import Field, ValidationError, validator
 
 from aind_data_schema.base import AindCoreModel, AindModel, BaseName, BaseNameEnumMeta, PIDName, Registry
 
@@ -271,9 +271,14 @@ class DataDescription(AindCoreModel):
     )
     data_summary: Optional[str] = Field(None, title="Data summary", description="Semantic summary of experimental goal")
 
+    # TODO: We need to remove all the custom class constructors on pydantic
+    #  models
     def __init__(self, label=None, **kwargs):
         """Construct a generic DataDescription"""
 
+        # Ideally, we'd like to just use validators to parse information,
+        # but we need to get rid of these init methods first since they
+        # don't get called on here
         super().__init__(**kwargs)
 
         if label is not None:
@@ -297,6 +302,23 @@ def parse_name(cls, name):
             creation_time=creation_time,
         )
 
+    @validator("data_level", pre=True, always=True)
+    def upgrade_data_level(cls, value: Union[str, DataLevel]):
+        """Updates legacy values to current values"""
+        # If user inputs a string and is 'raw level', convert it to RAW
+        if isinstance(value, str) and value in ["raw level", "raw data"]:
+            return DataLevel.RAW
+        # If user inputs a string, try to convert it to a DataLevel. Will raise
+        # an error if unable to parse the input string
+        elif isinstance(value, str):
+            return DataLevel(value)
+        # If user inputs a DataLevel object, return the object without parsing
+        elif isinstance(value, DataLevel):
+            return value
+        # else raise a validation error
+        else:
+            raise ValidationError("Data Level needs to be string or enum")
+
 
 class DerivedDataDescription(DataDescription):
     """A logical collection of data files derived via processing"""

src/aind_data_schema/schema_upgrade/data_description_upgrade.py

@@ -131,6 +131,7 @@ def upgrade_data_description(self, **kwargs) -> DataDescription:
             modality = [ModalityUpgrade.upgrade_modality(m) for m in old_modality]
         else:
             modality = getattr(DataDescription.__fields__.get("modality"), "default")
+        old_data_level = self._get_or_default(self.old_data_description_model, "data_level", kwargs)
 
         experiment_type = self._get_or_default(self.old_data_description_model, "experiment_type", kwargs)
         platform = None
@@ -155,7 +156,7 @@ def upgrade_data_description(self, **kwargs) -> DataDescription:
             name=self._get_or_default(self.old_data_description_model, "name", kwargs),
             institution=institution,
             funding_source=funding_source,
-            data_level=self._get_or_default(self.old_data_description_model, "data_level", kwargs),
+            data_level=old_data_level,
             group=self._get_or_default(self.old_data_description_model, "group", kwargs),
             investigators=self._get_or_default(self.old_data_description_model, "investigators", kwargs),
             project_name=self._get_or_default(self.old_data_description_model, "project_name", kwargs),

tests/resources/ephys_data_description/data_description_0.3.0_wrong_field.json

@@ -0,0 +1,17 @@
+{
+   "describedBy": "https://raw.githubusercontent.com/AllenNeuralDynamics/aind-data-schema/main/src/aind_data_schema/data_description.py",
+   "schema_version": "0.3.0",
+   "license": "CC-BY-4.0",
+   "creation_time": "10:52:15",
+   "creation_date": "2022-07-26",
+   "name": "ecephys_624643_2022-07-26_10-52-15",
+   "institution": "AIND",
+   "funding_source": [],
+   "data_level": "raw data",
+   "group": null,
+   "project_name": null,
+   "project_id": null,
+   "restrictions": null,
+   "modality": "ecephys",
+   "subject_id": "624643"
+}

tests/test_schema_upgrade.py

@@ -76,6 +76,79 @@ def test_upgrades_0_3_0(self):
         self.assertEqual([], new_data_description.related_data)
         self.assertIsNone(new_data_description.data_summary)
 
+    def test_upgrades_0_3_0_wrong_field(self):
+        """Tests data_description_0.3.0_wrong_field.json is mapped correctly."""
+        data_description_0_3_0 = self.data_descriptions["data_description_0.3.0_wrong_field.json"]
+        upgrader = DataDescriptionUpgrade(old_data_description_model=data_description_0_3_0)
+        # Should complain about platform being None and missing data level
+        with self.assertRaises(Exception) as e:
+            upgrader.upgrade_data_description()
+
+        expected_error_message = (
+            "ValidationError("
+            "model='DataDescription', "
+            "errors=[{"
+            "'loc': ('platform',), "
+            "'msg': 'none is not an allowed value', "
+            "'type': 'type_error.none.not_allowed'"
+            "}])"
+        )
+        self.assertEqual(expected_error_message, repr(e.exception))
+
+        # Should work by setting platform explicitly and DataLevel
+        new_data_description = upgrader.upgrade_data_description(platform=Platform.ECEPHYS, data_level=DataLevel.RAW)
+        self.assertEqual(datetime.datetime(2022, 7, 26, 10, 52, 15), new_data_description.creation_time)
+        self.assertEqual("ecephys_624643_2022-07-26_10-52-15", new_data_description.name)
+        self.assertEqual(Institution.AIND, new_data_description.institution)
+        self.assertEqual([], new_data_description.funding_source)
+        self.assertEqual(DataLevel.RAW, new_data_description.data_level)
+        self.assertIsNone(new_data_description.group)
+        self.assertEqual([], new_data_description.investigators)
+        self.assertIsNone(new_data_description.project_name)
+        self.assertIsNone(new_data_description.restrictions)
+        self.assertEqual([Modality.ECEPHYS], new_data_description.modality)
+        self.assertEqual("624643", new_data_description.subject_id)
+        self.assertEqual([], new_data_description.related_data)
+        self.assertIsNone(new_data_description.data_summary)
+
+        # Should also work by inputting legacy
+        new_data_description2 = upgrader.upgrade_data_description(platform=Platform.ECEPHYS, data_level="raw level")
+        self.assertEqual(DataLevel.RAW, new_data_description2.data_level)
+
+        # Should fail if inputting unknown string
+        with self.assertRaises(Exception) as e1:
+            upgrader.upgrade_data_description(platform=Platform.ECEPHYS, data_level="asfnewnjfq")
+
+        expected_error_message1 = (
+            "ValidationError(model='DataDescription', "
+            "errors=[{'loc': ('data_level',), "
+            "'msg': \"'asfnewnjfq' is not a valid DataLevel\", "
+            "'type': 'value_error'}])"
+        )
+
+        self.assertEqual(expected_error_message1, repr(e1.exception))
+
+        # Should also fail if inputting wrong type
+        with self.assertRaises(Exception) as e2:
+            upgrader.upgrade_data_description(platform=Platform.ECEPHYS, data_level=["raw"])
+        expected_error_message2 = (
+            "ValidationError(model='DataDescription', "
+            "errors=[{'loc': ('data_level',), "
+            "'msg': '__init__() takes exactly 3 positional arguments "
+            "(2 given)', 'type': 'type_error'}])"
+        )
+
+        self.assertEqual(expected_error_message2, repr(e2.exception))
+
+        # Should work if data_level is missing in original json doc and
+        # user sets it explicitly
+        data_description_dict = data_description_0_3_0.dict()
+        del data_description_dict["data_level"]
+        data_description_0_3_0_no_data_level = DataDescription.construct(**data_description_dict)
+        upgrader3 = DataDescriptionUpgrade(old_data_description_model=data_description_0_3_0_no_data_level)
+        new_data_description3 = upgrader3.upgrade_data_description(platform=Platform.ECEPHYS, data_level=DataLevel.RAW)
+        self.assertEqual(DataLevel.RAW, new_data_description3.data_level)
+
     def test_upgrades_0_4_0(self):
         """Tests data_description_0.4.0.json is mapped correctly."""
         data_description_0_4_0 = self.data_descriptions["data_description_0.4.0.json"]