Deduplicate FileDatastore functions

Remove a small amount of duplicated code by simplifying the function parameters
lsst · Nov 29, 2023 · 4c72aa3 · 4c72aa3
1 parent 199596b
commit 4c72aa3
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 15 deletions.
diff --git a/python/lsst/daf/butler/datastores/fileDatastore.py b/python/lsst/daf/butler/datastores/fileDatastore.py
@@ -731,14 +731,11 @@ def _prepare_for_direct_get(
                 if not fileLocations:
                     raise FileNotFoundError(f"None of the component files for dataset {ref} exist.")
 
-        # Is this a component request?
-        refComponent = ref.datasetType.component()
         return generate_datastore_get_information(
             fileLocations,
-            refStorageClass=refStorageClass,
-            refComponent=refComponent,
+            readStorageClass=refStorageClass,
+            ref=ref,
             parameters=parameters,
-            dataId=ref.dataId,
         )
 
     def _prepare_for_put(self, inMemoryDataset: Any, ref: DatasetRef) -> tuple[Location, Formatter]:

diff --git a/python/lsst/daf/butler/datastores/fileDatastoreClient.py b/python/lsst/daf/butler/datastores/fileDatastoreClient.py
@@ -48,10 +48,8 @@ def get_dataset_as_python_object(
 
     datastore_file_info = generate_datastore_get_information(
         fileLocations,
-        refStorageClass=ref.datasetType.storageClass,
-        refComponent=ref.datasetType.component(),
+        ref=ref,
         parameters=parameters,
-        dataId=ref.dataId,
     )
     return get_dataset_as_python_object_from_get_info(
         datastore_file_info, ref=ref, parameters=parameters, cache_manager=DatastoreDisabledCacheManager()

diff --git a/python/lsst/daf/butler/datastores/file_datastore/get.py b/python/lsst/daf/butler/datastores/file_datastore/get.py
@@ -36,7 +36,7 @@
 from dataclasses import dataclass
 from typing import Any, TypeAlias
 
-from lsst.daf.butler import DataCoordinate, DatasetRef, FileDescriptor, Formatter, Location, StorageClass
+from lsst.daf.butler import DatasetRef, FileDescriptor, Formatter, Location, StorageClass
 from lsst.daf.butler.datastore.cache_manager import AbstractDatastoreCacheManager
 from lsst.daf.butler.datastore.generic_base import post_process_get
 from lsst.daf.butler.datastore.stored_file_info import StoredFileInfo
@@ -80,11 +80,16 @@ class DatastoreFileGetInformation:
 def generate_datastore_get_information(
     fileLocations: list[DatasetLocationInformation],
     *,
-    refStorageClass: StorageClass,
-    refComponent: str | None,
+    ref: DatasetRef,
     parameters: Mapping[str, Any] | None,
-    dataId: DataCoordinate,
+    readStorageClass: StorageClass | None = None,
 ) -> list[DatastoreFileGetInformation]:
+    if readStorageClass is None:
+        readStorageClass = ref.datasetType.storageClass
+
+    # Is this a component request?
+    refComponent = ref.datasetType.component()
+
     disassembled = len(fileLocations) > 1
     fileGetInfo = []
     for location, storedFileInfo in fileLocations:
@@ -94,8 +99,6 @@ def generate_datastore_get_information(
         # If this has been disassembled we need read to match the write
         if disassembled:
             readStorageClass = writeStorageClass
-        else:
-            readStorageClass = refStorageClass
 
         formatter = get_instance_of(
             storedFileInfo.formatter,
@@ -105,7 +108,7 @@ def generate_datastore_get_information(
                 storageClass=writeStorageClass,
                 parameters=parameters,
             ),
-            dataId,
+            ref.dataId,
         )
 
         formatterParams, notFormatterParams = formatter.segregateParameters()