From ae4aec72cb030e33b638df7f55b1a6b18bc8d137 Mon Sep 17 00:00:00 2001
From: Christopher <chrism@alleninstitute.org>
Date: Mon, 15 Jul 2024 09:49:04 -0700
Subject: [PATCH] new attachments iteration

---
 src/aind_slims_api/core.py | 92 +++++++++++++++++++++-----------------
 1 file changed, 52 insertions(+), 40 deletions(-)

diff --git a/src/aind_slims_api/core.py b/src/aind_slims_api/core.py
index 8503294..ce57c02 100644
--- a/src/aind_slims_api/core.py
+++ b/src/aind_slims_api/core.py
@@ -28,6 +28,7 @@
 from slims.internal import (
     Column as SlimsColumn,
     Record as SlimsRecord,
+    Attachment as SlimsAttachment,
 )
 from slims.criteria import Criterion, conjunction, equals
 
@@ -94,7 +95,7 @@ class MyModel(SlimsBaseModel):
 
     pk: int = None
     json_entity: dict = None
-    attachments: Optional[Callable[[], Sequence['Attachment']]] = None
+    attachments: Optional[Callable[[], Sequence[SlimsAttachment]]] = None
     _slims_table: SLIMSTABLES
 
     @field_validator("*", mode="before")
@@ -142,45 +143,45 @@ def _serialize(self, field, info):
     # TODO: Add links - need Record.json_entity['links']['self']
     # TODO: Add Table - need Record.json_entity['tableName']
 
-    def fetch_attachments(self) -> Sequence['Attachment']:
-        """Fetches all attachments for this record"""
-        if not self.attachments:
-            logger.debug("Initialized without attachments.")
-            return []
-        validated = []
-        for attachment in self.attachments():
-            try:
-                validated.append(Attachment.model_validate(attachment))
-            except ValidationError as e:
-                logger.error(f"SLIMS data validation failed, {repr(e)}")
-        return validated
-
-    def fetch_attachments_content(self) -> \
-            Generator[Response, None, None]:
-        """Fetches the content of all attachments for this record and returns
-        them as dictionaries.
-
-        Notes
-        -----
-        - Assumes that the attachments are json
-        - Should this actually be text and not json?
-        """
-        for attachment in self.fetch_attachments():
-            yield attachment.fetch_content()
-
-
-class Attachment(SlimsBaseModel):
-
-    pk: int = Field(..., alias="attm_pk")
-    slims_api: Slims
-    _slims_table: SLIMSTABLES = "Attachment"
-
-    class Config:
-        arbitrary_types_allowed = True
-
-    def fetch_content(self) -> Response:
-        """Fetches the content of this attachment"""
-        return self.slims_api.get(f"repo/{self.pk}")
+    # def fetch_attachments(self) -> Sequence['Attachment']:
+    #     """Fetches all attachments for this record"""
+    #     if not self.attachments:
+    #         logger.debug("Initialized without attachments.")
+    #         return []
+    #     validated = []
+    #     for attachment in self.attachments():
+    #         try:
+    #             validated.append(Attachment.model_validate(attachment))
+    #         except ValidationError as e:
+    #             logger.error(f"SLIMS data validation failed, {repr(e)}")
+    #     return validated
+
+    # def fetch_attachments_content(self) -> \
+    #         Generator[Response, None, None]:
+    #     """Fetches the content of all attachments for this record and returns
+    #     them as dictionaries.
+
+    #     Notes
+    #     -----
+    #     - Assumes that the attachments are json
+    #     - Should this actually be text and not json?
+    #     """
+    #     for attachment in self.fetch_attachments():
+    #         yield attachment.fetch_content()
+
+
+# class Attachment(SlimsBaseModel):
+
+#     pk: int = Field(..., alias="attm_pk")
+#     slims_api: Slims
+#     _slims_table: SLIMSTABLES = "Attachment"
+
+#     class Config:
+#         arbitrary_types_allowed = True
+
+#     def fetch_content(self) -> Response:
+#         """Fetches the content of this attachment"""
+#         return self.slims_api.get(f"repo/{self.pk}")
 
     # attm_pk: SlimsColumn
     # attm_fk_content: SlimsColumn
@@ -353,6 +354,17 @@ def fetch_model(
             logger.debug(f"Found {len(records)} records for {model}.")
         return records[0]
 
+    def fetch_attachments_contents(
+        self,
+        record: SlimsBaseModel,
+    ) -> Generator[Response, None, None]:
+        """Fetches all attachments for a record"""
+        if not record.attachments:
+            raise ValueError("Record initialized no attachments.")
+
+        for attachment in record.attachments():
+            yield attachment.slims_api.get(f"repo/{attachment.attm_pk}")
+
     @lru_cache(maxsize=None)
     def fetch_pk(self, table: SLIMSTABLES, *args, **kwargs) -> int | None:
         """SlimsClient.fetch but returns the pk of the first returned record"""