feat: add pydantic model for slims records (#12)

* feat: add pydantic model for slims records

* fix: f-string typo

* feat: add unit model

* feat: include json_entity in models

* feat: return dict if validation fails

* test: update unit tests

* test: clear existing env vars in config test

* style: clean up core/user test

* test: add unitspec/mouse validation tests

* fix: return json_entity if validation fails

* style: change relative imports to absolute

src/aind_slims_api/core.py

@@ -8,13 +8,24 @@
     methods and integration with SlimsBaseModel subtypes
 """
 
-import logging
+from datetime import datetime
 from functools import lru_cache
+from pydantic import (
+    BaseModel,
+    ValidationInfo,
+    field_serializer,
+    field_validator,
+)
+from pydantic.fields import FieldInfo
+import logging
 from typing import Literal, Optional
 
-from slims.criteria import Criterion, conjunction, equals
-from slims.internal import Record as SlimsRecord
 from slims.slims import Slims, _SlimsApiException
+from slims.internal import (
+    Column as SlimsColumn,
+    Record as SlimsRecord,
+)
+from slims.criteria import Criterion, conjunction, equals
 
 from aind_slims_api import config
 
@@ -33,6 +44,94 @@
 ]
 
 
+class UnitSpec:
+    """Used in type annotation metadata to specify units"""
+
+    units: list[str]
+    preferred_unit: str = None
+
+    def __init__(self, *args, preferred_unit=None):
+        """Set list of acceptable units from args, and preferred_unit"""
+        self.units = args
+        if len(self.units) == 0:
+            raise ValueError("One or more units must be specified")
+        if preferred_unit is None:
+            self.preferred_unit = self.units[0]
+
+
+def _find_unit_spec(field: FieldInfo) -> UnitSpec | None:
+    """Given a Pydantic FieldInfo, find the UnitSpec in its metadata"""
+    metadata = field.metadata
+    for m in metadata:
+        if isinstance(m, UnitSpec):
+            return m
+    return None
+
+
+class SlimsBaseModel(
+    BaseModel,
+    from_attributes=True,
+    validate_assignment=True,
+):
+    """Pydantic model to represent a SLIMS record.
+    Subclass with fields matching those in the SLIMS record.
+
+    For Quantities, specify acceptable units like so:
+
+        class MyModel(SlimsBaseModel):
+            myfield: Annotated[float | None, UnitSpec("g","kg")]
+
+        Quantities will be serialized using the first unit passed
+
+    Datetime fields will be serialized to an integer ms timestamp
+    """
+
+    pk: int = None
+    json_entity: dict = None
+    _slims_table: SLIMSTABLES
+
+    @field_validator("*", mode="before")
+    def _validate(cls, value, info: ValidationInfo):
+        """Validates a field, accounts for Quantities"""
+        if isinstance(value, SlimsColumn):
+            if value.datatype == "QUANTITY":
+                unit_spec = _find_unit_spec(cls.model_fields[info.field_name])
+                if unit_spec is None:
+                    msg = (
+                        f'Quantity field "{info.field_name}"'
+                        "must be annotated with a UnitSpec"
+                    )
+                    raise TypeError(msg)
+                if value.unit not in unit_spec.units:
+                    msg = (
+                        f'Unexpected unit "{value.unit}" for field '
+                        f"{info.field_name}, Expected {unit_spec.units}"
+                    )
+                    raise ValueError(msg)
+            return value.value
+        else:
+            return value
+
+    @field_serializer("*")
+    def _serialize(self, field, info):
+        """Serialize a field, accounts for Quantities and datetime"""
+        unit_spec = _find_unit_spec(self.model_fields[info.field_name])
+        if unit_spec and field is not None:
+            quantity = {
+                "amount": field,
+                "unit_display": unit_spec.preferred_unit,
+            }
+            return quantity
+        elif isinstance(field, datetime):
+            return int(field.timestamp() * 10**3)
+        else:
+            return field
+
+    # TODO: Add links - need Record.json_entity['links']['self']
+    # TODO: Add Table - need Record.json_entity['tableName']
+    # TODO: Support attachments
+
+
 class SlimsClient:
     """Wrapper around slims-python-api client with convenience methods"""
 
@@ -133,3 +232,52 @@ def rest_link(self, table: SLIMSTABLES, **kwargs):
         base_url = f"{self.url}/rest/{table}"
         queries = [f"?{k}={v}" for k, v in kwargs.items()]
         return base_url + "".join(queries)
+
+    def add_model(self, model: SlimsBaseModel, *args, **kwargs) -> SlimsBaseModel:
+        """Given a SlimsBaseModel object, add it to SLIMS
+        Args
+            model (SlimsBaseModel): object to add
+            *args (str): fields to include in the serialization
+            **kwargs: passed to model.model_dump()
+
+        Returns
+            An instance of the same type of model, with data from
+            the resulting SLIMS record
+        """
+        fields_to_include = set(args) or None
+        fields_to_exclude = set(kwargs.get("exclude", []))
+        fields_to_exclude.add("pk")
+        rtn = self.add(
+            model._slims_table,
+            model.model_dump(
+                include=fields_to_include,
+                exclude=fields_to_exclude,
+                **kwargs,
+                by_alias=True,
+            ),
+        )
+        return type(model).model_validate(rtn)
+
+    def update_model(self, model: SlimsBaseModel, *args, **kwargs):
+        """Given a SlimsBaseModel object, update its (existing) SLIMS record
+
+        Args
+            model (SlimsBaseModel): object to update
+            *args (str): fields to include in the serialization
+            **kwargs: passed to model.model_dump()
+
+        Returns
+            An instance of the same type of model, with data from
+            the resulting SLIMS record
+        """
+        fields_to_include = set(args) or None
+        rtn = self.update(
+            model._slims_table,
+            model.pk,
+            model.model_dump(
+                include=fields_to_include,
+                by_alias=True,
+                **kwargs,
+            ),
+        )
+        return type(model).model_validate(rtn)

src/aind_slims_api/mouse.py

@@ -1,17 +1,51 @@
 """Contains a model for the mouse content, and a method for fetching it"""
 
 import logging
-from typing import Optional
+from typing import Annotated
 
-from aind_slims_api.core import SlimsClient
+from pydantic import Field, BeforeValidator, ValidationError
+
+from aind_slims_api.core import SlimsBaseModel, SlimsClient, UnitSpec, SLIMSTABLES
 
 logger = logging.getLogger()
 
 
+class SlimsMouseContent(SlimsBaseModel):
+    """Model for an instance of the Mouse ContentType"""
+
+    baseline_weight_g: Annotated[float | None, UnitSpec("g")] = Field(
+        ..., alias="cntn_cf_baselineWeight"
+    )
+    point_of_contact: str | None = Field(..., alias="cntn_cf_scientificPointOfContact")
+    water_restricted: Annotated[bool, BeforeValidator(lambda x: x or False)] = Field(
+        ..., alias="cntn_cf_waterRestricted"
+    )
+    barcode: str = Field(..., alias="cntn_barCode")
+    pk: int = Field(..., alias="cntn_pk")
+
+    _slims_table: SLIMSTABLES = "Content"
+
+    # TODO: Include other helpful fields (genotype, gender...)
+
+    # pk: callable
+    # cntn_fk_category: SlimsColumn
+    # cntn_fk_contentType: SlimsColumn
+    # cntn_barCode: SlimsColumn
+    # cntn_id: SlimsColumn
+    # cntn_cf_contactPerson: SlimsColumn
+    # cntn_status: SlimsColumn
+    # cntn_fk_status: SlimsColumn
+    # cntn_fk_user: SlimsColumn
+    # cntn_cf_fk_fundingCode: SlimsColumn
+    # cntn_cf_genotype: SlimsColumn
+    # cntn_cf_labtracksId: SlimsColumn
+    # cntn_cf_parentBarcode: SlimsColumn
+
+
 def fetch_mouse_content(
     client: SlimsClient,
     mouse_name: str,
-) -> Optional[dict]:
+) -> SlimsMouseContent | dict | None:
     """Fetches mouse information for a mouse with labtracks id {mouse_name}"""
     mice = client.fetch(
         "Content",
@@ -28,6 +62,12 @@ def fetch_mouse_content(
             )
     else:
         logger.warning("Warning, Mouse not in SLIMS")
-        mouse_details = None
+        return
+
+    try:
+        mouse = SlimsMouseContent.model_validate(mouse_details)
+    except ValidationError as e:
+        logger.error(f"SLIMS data validation failed, {repr(e)}")
+        return mouse_details.json_entity
 
-    return None if mouse_details is None else mouse_details.json_entity
+    return mouse

src/aind_slims_api/unit.py

@@ -0,0 +1,20 @@
+"""Contains a model for a unit"""
+
+import logging
+from typing import Optional
+
+from pydantic import Field
+
+from aind_slims_api.core import SlimsBaseModel
+
+logger = logging.getLogger()
+
+
+class SlimsUnit(SlimsBaseModel):
+    """Model for unit information in SLIMS"""
+
+    name: str = Field(..., alias="unit_name")
+    abbreviation: Optional[str] = Field("", alias="unit_abbreviation")
+    pk: int = Field(..., alias="unit_pk")
+
+    _slims_table: str = "Unit"

src/aind_slims_api/user.py

@@ -3,15 +3,31 @@
 import logging
 from typing import Optional
 
-from aind_slims_api.core import SlimsClient
+from pydantic import Field, ValidationError
+
+from aind_slims_api.core import SlimsBaseModel, SlimsClient
 
 logger = logging.getLogger()
 
 
+# TODO: Tighten this up once users are more commonly used
+class SlimsUser(SlimsBaseModel):
+    """Model for user information in SLIMS"""
+
+    username: str = Field(..., alias="user_userName")
+    first_name: Optional[str] = Field("", alias="user_firstName")
+    last_name: Optional[str] = Field("", alias="user_lastName")
+    full_name: Optional[str] = Field("", alias="user_fullName")
+    email: Optional[str] = Field("", alias="user_email")
+    pk: int = Field(..., alias="user_pk")
+
+    _slims_table: str = "User"
+
+
 def fetch_user(
     client: SlimsClient,
     username: str,
-) -> Optional[dict]:
+) -> SlimsUser | dict | None:
     """Fetches user information for a user with username {username}"""
     users = client.fetch(
         "User",
@@ -23,11 +39,17 @@ def fetch_user(
         if len(users) > 1:
             logger.warning(
                 f"Warning, Multiple users in SLIMS with "
-                f"username {[u.json_entity for u in users]}, "
+                f"username {username}, "
                 f"using pk={user_details.pk()}"
             )
     else:
         logger.warning("Warning, User not in SLIMS")
-        user_details = None
+        return
+
+    try:
+        user = SlimsUser.model_validate(user_details)
+    except ValidationError as e:
+        logger.error(f"SLIMS data validation failed, {repr(e)}")
+        return user_details.json_entity
 
-    return None if user_details is None else user_details.json_entity
+    return user

tests/test_configuration.py

@@ -10,6 +10,11 @@
 class TestAindSlimsApiSettings(unittest.TestCase):
     """Tests methods in AindSlimsApiSettings class"""
 
+    @patch.dict(
+        os.environ,
+        {},
+        clear=True,
+    )
     def test_default_settings(self):
         """Tests that the class will be set with defaults"""
         default_settings = AindSlimsApiSettings()