Fix logic bug in _getStorageClass and add support for str type

lsst · Mar 15, 2024 · 00859ee · 00859ee
1 parent 511d1ef
commit 00859ee
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/python/lsst/pipe/base/_dataset_handle.py b/python/lsst/pipe/base/_dataset_handle.py
@@ -79,7 +79,7 @@ def __init__(
         self,
         inMemoryDataset: Any,
         *,
-        storageClass: StorageClass | None = None,
+        storageClass: StorageClass | str | None = None,
         parameters: dict[str, Any] | None = None,
         dataId: DataId | None = None,
         copy: bool = False,
@@ -252,7 +252,10 @@ def _getStorageClass(self) -> StorageClass:
         """
         factory = StorageClassFactory()
         if self.storageClass:
-            return factory.getStorageClass(self.storageClass)
+            if isinstance(self.storageClass, str):
+                return factory.getStorageClass(self.storageClass)
+            else:
+                return self.storageClass
 
         # Need to match python type.
         pytype = type(self.inMemoryDataset)