Merge pull request #37 from AllenNeuralDynamics/feat-reencode-aind-data-schema-models

Add encoder for special characters for aind-data-schema models

Author: bruno-f-cruz

src/aind_behavior_experiment_launcher/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.3.1"
+__version__ = "0.3.2"
 
 import logging
 import logging.config

src/aind_behavior_experiment_launcher/data_mapper/aind_data_schema.py

@@ -1,37 +1,78 @@
 from __future__ import annotations
 
-try:
-    import aind_data_schema
-except ImportError as _e:
-    _e.add_note(
-        "The 'aind-data-schema' package is required to use this module. \
-            Install the optional dependencies defined in `project.toml' \
-                by running `pip install .[aind-services]`"
+import importlib.util
+
+if importlib.util.find_spec("aind_data_schema") is None:
+    raise ImportError(
+        "The 'aind-data-schema' package is required to use this module. "
+        "Install the optional dependencies defined in `project.toml` "
+        "by running `pip install .[aind-services]`"
     )
-    raise _e
 
 import abc
 import logging
-from typing import Generic, TypeVar, Union
-
-import aind_data_schema.components.devices
-import aind_data_schema.core.rig
-import aind_data_schema.core.session
+from typing import Any, Generic, Type, TypeVar, Union
 
-from . import _base
+from aind_data_schema.core import rig as ads_rig
+from aind_data_schema.core import session as ads_session
+from pydantic import BaseModel, create_model, model_validator
 
-TAdsObject = TypeVar("TAdsObject", bound=Union[aind_data_schema.core.session.Session, aind_data_schema.core.rig.Rig])
+from aind_behavior_experiment_launcher.data_mapper import _base
 
 logger = logging.getLogger(__name__)
 
+_TAdsObject = TypeVar("_TAdsObject", bound=Union[ads_session.Session, ads_rig.Rig])
 
-class AindDataSchemaDataMapper(_base.DataMapper[TAdsObject], abc.ABC, Generic[TAdsObject]):
+
+class AindDataSchemaDataMapper(_base.DataMapper[_TAdsObject], abc.ABC, Generic[_TAdsObject]):
     @property
     @abc.abstractmethod
     def session_name(self) -> str: ...
 
 
-class AindDataSchemaSessionDataMapper(AindDataSchemaDataMapper[aind_data_schema.core.session.Session], abc.ABC): ...
+class AindDataSchemaSessionDataMapper(AindDataSchemaDataMapper[ads_session.Session], abc.ABC): ...
+
+
+class AindDataSchemaRigDataMapper(AindDataSchemaDataMapper[ads_rig.Rig], abc.ABC): ...
+
+
+_TModel = TypeVar("_TModel", bound=BaseModel)
 
 
-class AindDataSchemaRigDataMapper(AindDataSchemaDataMapper[aind_data_schema.core.rig.Rig], abc.ABC): ...
+def create_encoding_model(model: Type[_TModel]) -> Type[_TModel]:
+    """Creates a new BaseModel by wrapping the incoming model and adding a Before
+    ModelValidator to replace _SPECIAL_CHARACTERS with the unicode, escaped,
+    representation"""
+
+    _SPECIAL_CHARACTERS = ".$"
+
+    def _to_unicode_repr(character: str):
+        if len(character) != 1:
+            raise ValueError(f"Expected a single character, got {character}")
+        return f"\\u{ord(character):04x}"
+
+    def _aind_data_schema_encoder(cls, data: Any) -> Any:
+        if isinstance(data, dict):
+            return _sanitize_dict(data)
+        return data
+
+    def _sanitize_dict(value: dict) -> dict:
+        if isinstance(value, dict):
+            _keys = list(value.keys())
+            for key in _keys:
+                if isinstance(value[key], dict):
+                    value[key] = _sanitize_dict(value[key])
+                if isinstance(sanitized_key := key, str):
+                    for char in _SPECIAL_CHARACTERS:
+                        if char in sanitized_key:
+                            sanitized_key = sanitized_key.replace(char, _to_unicode_repr(char))
+                    value[sanitized_key] = value.pop(key)
+        return value
+
+    return create_model(
+        f"_Wrapped{model.__class__.__name__}",
+        __base__=model,
+        __validators__={
+            "encoder": model_validator(mode="before")(_aind_data_schema_encoder)  # type: ignore
+        },
+    )

tests/test_data_mapper.py

@@ -1,10 +1,12 @@
 import unittest
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Any, Dict, List, Optional
 from unittest.mock import patch
 
+from aind_data_schema.base import AindGeneric
 from pydantic import BaseModel
 
+from aind_behavior_experiment_launcher.data_mapper.aind_data_schema import create_encoding_model
 from aind_behavior_experiment_launcher.data_mapper.helpers import (
     _sanity_snapshot_keys,
     snapshot_bonsai_environment,
@@ -65,5 +67,19 @@ def test_sanity_snapshot_keys_with_dots_and_dollars(self):
         self.assertEqual(result, expected)
 
 
+class TestAindDataMapper(unittest.TestCase):
+    class MyMockModel(BaseModel):
+        a_dict: Dict[str, Any]
+        a_generic: AindGeneric
+
+    def test_encoding_with_illegal_characters(self):
+        _input = {"key": "value", "$key.key": "value"}
+        _expected = {"key": "value", "\\u0024key\\u002ekey": "value"}
+        encoding_model = create_encoding_model(self.MyMockModel)
+        test = encoding_model(a_dict=_input, a_generic=_input)
+        self.assertEqual(test.a_dict, _expected)
+        self.assertEqual(test.a_generic.model_dump(), _expected)
+
+
 if __name__ == "__main__":
     unittest.main()