fix: add identified_objects to _models.py (#2015)

* fix: add identified_objects to _models.py

* add tests for identifiedobject _models.py

Author: daniel-ji

client/python/cryoet_data_portal/src/cryoet_data_portal/_codegen.py

@@ -69,6 +69,7 @@
     "Frame": "Frame",
     "FrameAcquisitionFile": "FrameAcquisitionFile",
     "GainFile": "GainFile",
+    "IdentifiedObject": "IdentifiedObject",
     "PerSectionAlignmentParameters": "PerSectionAlignmentParameters",
     "PerSectionParameters": "PerSectionParameters",
     "Run": "Run",

client/python/cryoet_data_portal/src/cryoet_data_portal/_models.py

@@ -1195,6 +1195,64 @@ def get_by_id(cls, client: Client, id: int):
     )
 
 
+class IdentifiedObject(Model):
+    """An ontology term/ID or accession code, state and description associated with a Run, describing cellular components that were visually identified, but not labeled, in a tomogram associated with the Run
+
+    Attributes:
+        id (int): Numeric identifier (May change!)
+        run (Run): The run this identified object is a part of
+        run_id (int): None
+        object_id (str): Gene Ontology Cellular Component identifier or UniProtKB accession for the identified object.
+        object_name (str): Name of the identified object (e.g. ribosome, nuclear pore complex, actin filament, membrane)
+        object_description (str): A textual description of the identified object, can be a longer description to include additional information not covered by the IdentifiedObject object name and state.
+        object_state (str): Molecule state of the identified object (e.g. open, closed)
+    """
+
+    _gql_type: str = "IdentifiedObject"
+    _gql_root_field: str = "identifiedObjects"
+
+    id: int = IntField()
+    run: Run = ItemRelationship("Run", "run_id", "id")
+    run_id: int = IntField()
+    object_id: str = StringField()
+    object_name: str = StringField()
+    object_description: str = StringField()
+    object_state: str = StringField()
+
+    @classmethod
+    def find(
+        cls,
+        client: Client,
+        query_filters: Optional[Iterable[GQLExpression]] = None,
+    ):
+        """
+        Examples:
+            >>> identified_objects = IdentifiedObject.find(client, query_filters=[IdentifiedObject.run.name == "TS_026"])
+
+            Get all results for this type:
+
+            >>> identified_objects = IdentifiedObject.find(client)
+        """
+        return super(IdentifiedObject, cls).find(client, query_filters)
+
+    find.__func__.__doc__ = Model.find.__func__.__doc__ + find.__func__.__doc__
+
+    @classmethod
+    def get_by_id(cls, client: Client, id: int):
+        """
+        Examples:
+            Get an IdentifiedObject by ID:
+
+            >>> identified_object = IdentifiedObject.get_by_id(client, 10000)
+            >>> print(identified_object.id)
+        """
+        return super(IdentifiedObject, cls).get_by_id(client, id)
+
+    get_by_id.__func__.__doc__ = (
+        Model.get_by_id.__func__.__doc__ + get_by_id.__func__.__doc__
+    )
+
+
 class PerSectionAlignmentParameters(Model):
     """Map alignment parameters to tilt series frames
 
@@ -1339,6 +1397,7 @@ class Run(Model):
         dataset_id (int): None
         frames (List[Frame]): The frames of this run
         gain_files (List[GainFile]): The gain files of this run
+        identified_objects (List[IdentifiedObject]): The identified objects of this run
         frame_acquisition_files (List[FrameAcquisitionFile]): The frame acquisition files of this run
         per_section_parameters (List[PerSectionParameters]): The per section parameters of this run
         tiltseries (List[TiltSeries]): The tilt series of this run
@@ -1359,6 +1418,11 @@ class Run(Model):
     dataset_id: int = IntField()
     frames: List[Frame] = ListRelationship("Frame", "id", "run_id")
     gain_files: List[GainFile] = ListRelationship("GainFile", "id", "run_id")
+    identified_objects: List[IdentifiedObject] = ListRelationship(
+        "IdentifiedObject",
+        "id",
+        "run_id",
+    )
     frame_acquisition_files: List[FrameAcquisitionFile] = ListRelationship(
         "FrameAcquisitionFile",
         "id",
@@ -1951,6 +2015,7 @@ def download_everything(self, dest_path: Optional[str] = None):
 Frame.setup()
 FrameAcquisitionFile.setup()
 GainFile.setup()
+IdentifiedObject.setup()
 PerSectionAlignmentParameters.setup()
 PerSectionParameters.setup()
 Run.setup()

client/python/cryoet_data_portal/tests/test_filters.py

@@ -1,6 +1,7 @@
 import pytest
 
 from cryoet_data_portal import Run, Tomogram
+from cryoet_data_portal._models import IdentifiedObject
 
 
 def test_basic_filters(client) -> None:
@@ -38,3 +39,38 @@ def test_filter_on_object_raises_exceptions(client) -> None:
         Run.find(client, [Run.annotations == 20001])
     with pytest.raises(Exception, match=match):
         Run.find(client, [Run.dataset == 20001])
+
+
+def test_identified_object_find(client) -> None:
+    all_objects = IdentifiedObject.find(client)
+    assert len(all_objects) == 3
+
+    ribosomes = IdentifiedObject.find(
+        client,
+        [IdentifiedObject.object_name == "ribosome"],
+    )
+    assert len(ribosomes) == 2
+    for obj in ribosomes:
+        assert obj.object_name == "ribosome"
+
+    go_objects = IdentifiedObject.find(
+        client,
+        [IdentifiedObject.object_id == "GO:0005739"],
+    )
+    assert len(go_objects) == 1
+    assert go_objects[0].object_name == "mitochondrion"
+
+
+def test_identified_object_find_with_run_filter(client) -> None:
+    run1_objects = IdentifiedObject.find(
+        client,
+        [IdentifiedObject.run.name == "RUN1"],
+    )
+    assert len(run1_objects) == 2
+
+    run2_objects = IdentifiedObject.find(
+        client,
+        [IdentifiedObject.run.name == "RUN2"],
+    )
+    assert len(run2_objects) == 1
+    assert run2_objects[0].object_name == "ribosome"

client/python/cryoet_data_portal/tests/test_get_by_id.py

@@ -1,7 +1,18 @@
 from cryoet_data_portal import Dataset
 from cryoet_data_portal._client import Client
+from cryoet_data_portal._models import IdentifiedObject
 
 
 def test_get_by_id(client: Client) -> None:
     dataset = Dataset.get_by_id(client, 20001)
     assert dataset.id == 20001
+
+
+def test_identified_object_get_by_id(client: Client) -> None:
+    identified_object = IdentifiedObject.get_by_id(client, 30001)
+    assert identified_object.id == 30001
+    assert identified_object.object_id == "GO:0005840"
+    assert identified_object.object_name == "ribosome"
+    assert identified_object.object_description == "A large ribonucleoprotein complex"
+    assert identified_object.object_state == "assembled"
+    assert identified_object.run_id is not None

client/python/cryoet_data_portal/tests/test_infra/seed_db.py

@@ -16,6 +16,7 @@
 from test_infra.factories.dataset_funding import DatasetFundingFactory
 from test_infra.factories.deposition import DepositionFactory
 from test_infra.factories.deposition_author import DepositionAuthorFactory
+from test_infra.factories.identified_object import IdentifiedObjectFactory
 from test_infra.factories.run import RunFactory
 from test_infra.factories.tiltseries import TiltseriesFactory
 from test_infra.factories.tomogram import TomogramFactory
@@ -489,6 +490,31 @@ def use_factoryboy() -> None:
     TomogramAuthorFactory.create(tomogram=tomo4, author_list_order=1)
     TomogramAuthorFactory.create(tomogram=tomo4, author_list_order=2)
 
+    IdentifiedObjectFactory.create(
+        id=30001,
+        run=r1,
+        object_id="GO:0005840",
+        object_name="ribosome",
+        object_description="A large ribonucleoprotein complex",
+        object_state="assembled",
+    )
+    IdentifiedObjectFactory.create(
+        id=30002,
+        run=r1,
+        object_id="GO:0005739",
+        object_name="mitochondrion",
+        object_description="A semiautonomous organelle",
+        object_state=None,
+    )
+    IdentifiedObjectFactory.create(
+        id=30003,
+        run=r2,
+        object_id="GO:0005840",
+        object_name="ribosome",
+        object_description="A large ribonucleoprotein complex",
+        object_state="assembled",
+    )
+
     FileFactory.update_file_ids()
 
     session.commit()