From ac0cd43c208bca2ed7ec585a83253f3c9c9264ed Mon Sep 17 00:00:00 2001
From: Christopher <chrism@alleninstitute.org>
Date: Tue, 16 Jul 2024 12:33:49 -0700
Subject: [PATCH] new core

---
 src/aind_slims_api/behavior_session.py | 23 ++++++++++++-----------
 src/aind_slims_api/core.py             | 19 ++++++++++++++-----
 src/aind_slims_api/instrument.py       |  2 +-
 3 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/src/aind_slims_api/behavior_session.py b/src/aind_slims_api/behavior_session.py
index 259178e..d271caa 100644
--- a/src/aind_slims_api/behavior_session.py
+++ b/src/aind_slims_api/behavior_session.py
@@ -9,6 +9,9 @@
 from pydantic import Field
 
 from aind_slims_api.core import SlimsBaseModel, SlimsClient, SLIMSTABLES
+from aind_slims_api.mouse import SlimsMouseContent
+from aind_slims_api.instrument import SlimsInstrument
+from aind_slims_api.user import SlimsUser
 
 logger = logging.getLogger()
 
@@ -31,7 +34,8 @@ class SlimsBehaviorSessionContentEvent(SlimsBaseModel):
     task_schema_version: str | None = Field(
         default=None, alias="cnvn_cf_taskSchemaVersion"
     )
-    software_version: str | None = Field(default=None, alias="cnvn_cf_softwareVersion")
+    software_version: str | None = Field(
+        default=None, alias="cnvn_cf_softwareVersion")
     date: datetime | None = Field(..., alias="cnvn_cf_scheduledDate")
     cnvn_fk_contentEventType: int = 10  # pk of Behavior Session ContentEvent
     _slims_table: SLIMSTABLES = "ContentEvent"
@@ -66,17 +70,14 @@ def _resolve_pk(
 
 def fetch_behavior_session_content_events(
     client: SlimsClient,
-    mouse: SlimsSingletonFetchReturn,
-) -> tuple[list[SlimsBehaviorSessionContentEvent], list[dict[str, Any]]]:
+    mouse: SlimsMouseContent,
+) -> list[SlimsBehaviorSessionContentEvent]:
     """Fetches behavior sessions for a mouse with labtracks id {mouse_name}
 
     Returns
     -------
-    tuple:
-        list:
-            Validated SlimsBehaviorSessionContentEvent objects
-        list:
-            Dictionaries representations of objects that failed validation
+    list:
+        Validated SlimsBehaviorSessionContentEvent objects
     """
     return client.fetch_models(
         SlimsBehaviorSessionContentEvent,
@@ -88,9 +89,9 @@ def fetch_behavior_session_content_events(
 
 def write_behavior_session_content_events(
     client: SlimsClient,
-    mouse: SlimsSingletonFetchReturn,
-    instrument: SlimsSingletonFetchReturn,
-    trainers: list[SlimsSingletonFetchReturn],
+    mouse: SlimsMouseContent,
+    instrument: SlimsInstrument,
+    trainers: list[SlimsUser],
     *behavior_sessions: SlimsBehaviorSessionContentEvent,
 ) -> list[SlimsBehaviorSessionContentEvent]:
     """Writes behavior sessions for a mouse with labtracks id {mouse_name}
diff --git a/src/aind_slims_api/core.py b/src/aind_slims_api/core.py
index 11b7055..6754f73 100644
--- a/src/aind_slims_api/core.py
+++ b/src/aind_slims_api/core.py
@@ -18,10 +18,11 @@
     field_serializer,
     field_validator,
 )
+from copy import deepcopy
 from pydantic.fields import FieldInfo
 import logging
 from typing import (
-    Any, Generator, Callable, Literal, Optional, Sequence, Type, TypeVar
+    Any, Generator, Callable, Literal, Optional, Sequence, Type, TypeVar,
 )
 from requests import Response
 from slims.slims import Slims, _SlimsApiException
@@ -97,6 +98,7 @@ class MyModel(SlimsBaseModel):
     json_entity: dict = None
     attachments: Optional[Callable[[], Sequence[SlimsAttachment]]] = None
     _slims_table: SLIMSTABLES
+    _base_fetch_filters: dict[str, str | int]  # use for fetch_models, fetch_model
 
     @field_validator("*", mode="before")
     def _validate(cls, value, info: ValidationInfo):
@@ -310,10 +312,17 @@ def fetch_models(
         -----
         - kwargs are mapped to field alias values
         """
-        resolved_kwargs = {
-            self.resolve_model_alias(model, name): value
-            for name, value in kwargs.items()
-        }
+        resolved_kwargs = deepcopy(model._base_fetch_filters)
+        for name, value in kwargs.items():
+            resolved_kwargs[self.resolve_model_alias(model, name)] = value
+        # resolved_kwargs = {
+        #     self.resolve_model_alias(model, name): value
+        #     for name, value in kwargs.items()
+        # }
+        # for filter_name, filter_value in model._base_fetch_filters:
+        #     if filter_name in resolved_kwargs:
+
+        #     resolved_kwargs[filter_name] = filter_value
         logger.debug("Resolved kwargs: %s", resolved_kwargs)
         resolved_sort: Optional[str | list[str]] = None
         if sort is not None:
diff --git a/src/aind_slims_api/instrument.py b/src/aind_slims_api/instrument.py
index ed8aaed..10190dd 100644
--- a/src/aind_slims_api/instrument.py
+++ b/src/aind_slims_api/instrument.py
@@ -11,7 +11,7 @@
 
 
 class SlimsInstrument(SlimsBaseModel):
-    """Model for an instance of the Behavior Session ContentEvent
+    """Model for a SLIMS instrument record.
 
     Examples
     --------