DM-35815: Add StorageClassFactory.findStorageClass

.github/workflows/build.yaml

@@ -69,7 +69,7 @@ jobs:
         shell: bash -l {0}
         run: |
           conda install -y -q \
-            flake8 \
+            "flake8<5" \
             pytest pytest-flake8 pytest-xdist pytest-openfiles pytest-cov
 
       - name: List installed packages

doc/changes/DM-35815.feature.rst

@@ -0,0 +1 @@
+Add ``StorageClass.findStorageClass`` method to find a storage class from a python type.

python/lsst/daf/butler/core/storageClass.py

@@ -439,14 +439,18 @@ def validateInstance(self, instance: Any) -> bool:
         """
         return isinstance(instance, self.pytype)
 
-    def is_type(self, other: Type) -> bool:
+    def is_type(self, other: Type, compare_types: bool = False) -> bool:
         """Return Boolean indicating whether the supplied type matches
         the type in this `StorageClass`.
 
         Parameters
         ----------
         other : `Type`
             The type to be checked.
+        compare_types : `bool`, optional
+            If `True` the python type will be used in the comparison
+            if the type names do not match. This may trigger an import
+            of code and so can be slower.
 
         Returns
         -------
@@ -463,7 +467,17 @@ def is_type(self, other: Type) -> bool:
             return self._pytype is other
 
         other_name = get_full_type_name(other)
-        return self._pytypeName == other_name
+        if self._pytypeName == other_name:
+            return True
+
+        if compare_types:
+            # Must protect against the import failing.
+            try:
+                return self.pytype is other
+            except Exception:
+                pass
+
+        return False
 
     def can_convert(self, other: StorageClass) -> bool:
         """Return `True` if this storage class can convert python types
@@ -842,6 +856,73 @@ def getStorageClass(self, storageClassName: str) -> StorageClass:
         """
         return self._storageClasses[storageClassName]
 
+    def findStorageClass(self, pytype: Type, compare_types: bool = False) -> StorageClass:
+        """Find the storage class associated with this python type.
+
+        Parameters
+        ----------
+        pytype : `type`
+            The Python type to be matched.
+        compare_types : `bool`, optional
+            If `False`, the type will be checked against name of the python
+            type. This comparison is always done first. If `True` and the
+            string comparison failed, each candidate storage class will be
+            forced to have its type imported. This can be significantly slower.
+
+        Returns
+        -------
+        storageClass : `StorageClass`
+            The matching storage class.
+
+        Raises
+        ------
+        KeyError
+            Raised if no match could be found.
+
+        Notes
+        -----
+        It is possible for a python type to be associated with multiple
+        storage classes. This method will currently return the first that
+        matches.
+        """
+        result = self._find_storage_class(pytype, False)
+        if result:
+            return result
+
+        if compare_types:
+            # The fast comparison failed and we were asked to try the
+            # variant that might involve code imports.
+            result = self._find_storage_class(pytype, True)
+            if result:
+                return result
+
+        raise KeyError(f"Unable to find a StorageClass associated with type {get_full_type_name(pytype)!r}")
+
+    def _find_storage_class(self, pytype: Type, compare_types: bool) -> Optional[StorageClass]:
+        """Iterate through all storage classes to find a match.
+
+        Parameters
+        ----------
+        pytype : `type`
+            The Python type to be matched.
+        compare_types : `bool`, optional
+            Whether to use type name matching or explicit type matching.
+            The latter can be slower.
+
+        Returns
+        -------
+        storageClass : `StorageClass` or `None`
+            The matching storage class, or `None` if no match was found.
+
+        Notes
+        -----
+        Helper method for ``findStorageClass``.
+        """
+        for storageClass in self.values():
+            if storageClass.is_type(pytype, compare_types=compare_types):
+                return storageClass
+        return None
+
     def registerStorageClass(self, storageClass: StorageClass, msg: Optional[str] = None) -> None:
         """Store the `StorageClass` in the factory.
 

tests/test_storageClass.py

@@ -40,6 +40,18 @@ class PythonType:
     pass
 
 
+class PythonType2:
+    """A dummy class to test the registry of Python types."""
+
+    pass
+
+
+class PythonType3:
+    """A dummy class to test the registry of Python types."""
+
+    pass
+
+
 class StorageClassFactoryTestCase(unittest.TestCase):
     """Tests of the storage class infrastructure."""
 
@@ -206,6 +218,30 @@ def testRegistry(self):
         # Check you can silently insert something that is already there
         factory.registerStorageClass(newclass3)
 
+    def testFactoryFind(self):
+        # Finding a storage class can involve doing lots of slow imports so
+        # this is a separate test.
+        factory = StorageClassFactory()
+        className = "PythonType3"
+        newclass = StorageClass(className, pytype=PythonType3)
+        factory.registerStorageClass(newclass)
+        sc = factory.getStorageClass(className)
+
+        # Can we find a storage class from a type.
+        new_sc = factory.findStorageClass(PythonType3)
+        self.assertEqual(new_sc, sc)
+
+        # Now with slow mode
+        new_sc = factory.findStorageClass(PythonType3, compare_types=True)
+        self.assertEqual(new_sc, sc)
+
+        # This class will never match.
+        with self.assertRaises(KeyError):
+            factory.findStorageClass(PythonType2, compare_types=True)
+
+        # Check builtins.
+        self.assertEqual(factory.findStorageClass(dict), factory.getStorageClass("StructuredDataDict"))
+
     def testFactoryConfig(self):
         factory = StorageClassFactory()
         factory.addFromConfig(StorageClassConfig())