Remove use of model as input to cached function

AllenNeuralDynamics · Aug 19, 2024 · 96e73d0 · 96e73d0
1 parent 89a7c6f
commit 96e73d0
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/src/aind_session/extensions/ecephys.py b/src/aind_session/extensions/ecephys.py
@@ -155,7 +155,7 @@ def sorted_data_dir(self) -> upath.UPath:
             )
             sorted_data_dir = (
                 aind_session.utils.codeocean_utils.get_data_asset_source_dir(
-                    self.sorted_data_asset.id
+                    asset_id=self.sorted_data_asset.id
                 )
             )
             logger.debug(
@@ -206,7 +206,7 @@ def compressed_dir(self) -> upath.UPath:
 
     @staticmethod
     def get_clipped_and_compressed_dirs(
-        raw_data_asset_id: str | uuid.UUID | codeocean.data_asset.DataAsset
+        raw_data_asset_id: str | uuid.UUID
     ) -> tuple[upath.UPath | None, upath.UPath | None]:
         """
         Paths to the dirs containing Open Ephys recording data in CodeOcean upload
@@ -221,7 +221,7 @@ def get_clipped_and_compressed_dirs(
         >>> clipped.as_posix()
         's3://aind-ephys-data/ecephys_676909_2023-12-13_13-43-40/ecephys_clipped'
         """
-        raw_data_dir = aind_session.utils.get_data_asset_source_dir(raw_data_asset_id)
+        raw_data_dir = aind_session.utils.get_data_asset_source_dir(asset_id=raw_data_asset_id)
         candidate_parent_dirs = (
             raw_data_dir / "ecephys",  # newer location in dedicated modality folder
             raw_data_dir,  # original location in root if upload folder
@@ -258,10 +258,10 @@ def sorted_probes(self) -> tuple[str, ...]:
         >>> session.ecephys.sorted_probes
         ('ProbeA', 'ProbeB', 'ProbeC', 'ProbeD', 'ProbeE', 'ProbeF')
         """
-        return self.get_sorted_probe_names(self.sorted_data_asset)
+        return self.get_sorted_probe_names(self.sorted_data_asset.id)
 
     @staticmethod
-    def get_sorted_probe_names(sorted_data_asset_id: str | uuid.UUID | codeocean.data_asset.DataAsset) -> tuple[str, ...]:
+    def get_sorted_probe_names(sorted_data_asset_id: str | uuid.UUID) -> tuple[str, ...]:
         """Names of probes that reached the final stage of the sorting pipeline.
 
         - checks for probe dirs in the asset's data dir
@@ -276,7 +276,7 @@ def get_sorted_probe_names(sorted_data_asset_id: str | uuid.UUID | codeocean.dat
         >>> aind_session.ecephys.get_sorted_probe_names('a2a54575-b5ca-4cf0-acd0-2933e18bcb2d')
         ('ProbeA', 'ProbeB', 'ProbeC', 'ProbeD', 'ProbeE', 'ProbeF')
         """
-        sorted_data_dir = aind_session.utils.get_data_asset_source_dir(sorted_data_asset_id)
+        sorted_data_dir = aind_session.utils.get_data_asset_source_dir(asset_id=sorted_data_asset_id)
         candidate_parent_dirs = (
             sorted_data_dir / "curated",
             sorted_data_dir / "sorting_precurated",

diff --git a/src/aind_session/session.py b/src/aind_session/session.py
@@ -214,7 +214,7 @@ def raw_data_dir(self) -> upath.UPath:
                 f"Using asset {self.raw_data_asset.id} to find raw data path for {self.id}"
             )
             raw_data_dir = aind_session.utils.get_data_asset_source_dir(
-                asset_id=self.raw_data_asset,
+                asset_id=self.raw_data_asset.id,
             )
             logger.debug(f"Raw data dir found for {self.id}: {raw_data_dir}")
             return raw_data_dir