Merge pull request #719 from lsst/tickets/DM-35741

DM-35741: Add some improvements to StorageClass

doc/changes/DM-29835.feature.rst

@@ -0,0 +1,2 @@
+* Add ``StorageClass.is_type`` method to compare a type with that of the storage class itelf.
+* Add keys, values, items, and iterator for ``StorageClassFactory``.

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

@@ -29,7 +29,23 @@
 import copy
 import itertools
 import logging
-from typing import Any, Collection, Dict, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union
+from typing import (
+    Any,
+    Collection,
+    Dict,
+    ItemsView,
+    Iterator,
+    KeysView,
+    List,
+    Mapping,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Type,
+    Union,
+    ValuesView,
+)
 
 from lsst.utils import doImportType
 from lsst.utils.classes import Singleton
@@ -423,6 +439,32 @@ def validateInstance(self, instance: Any) -> bool:
         """
         return isinstance(instance, self.pytype)
 
+    def is_type(self, other: Type) -> bool:
+        """Return Boolean indicating whether the supplied type matches
+        the type in this `StorageClass`.
+
+        Parameters
+        ----------
+        other : `Type`
+            The type to be checked.
+
+        Returns
+        -------
+        match : `bool`
+            `True` if the types are equal.
+
+        Notes
+        -----
+        If this `StorageClass` has not yet imported the Python type the
+        check is done against the full type name, this prevents an attempt
+        to import the type when it will likely not match.
+        """
+        if self._pytype:
+            return self._pytype is other
+
+        other_name = get_full_type_name(other)
+        return self._pytypeName == other_name
+
     def can_convert(self, other: StorageClass) -> bool:
         """Return `True` if this storage class can convert python types
         in the other storage class.
@@ -645,6 +687,21 @@ def __contains__(self, storageClassOrName: Union[StorageClass, str]) -> bool:
                 return storageClassOrName == self._storageClasses[storageClassOrName.name]
         return False
 
+    def __len__(self) -> int:
+        return len(self._storageClasses)
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._storageClasses)
+
+    def values(self) -> ValuesView[StorageClass]:
+        return self._storageClasses.values()
+
+    def keys(self) -> KeysView[str]:
+        return self._storageClasses.keys()
+
+    def items(self) -> ItemsView[str, StorageClass]:
+        return self._storageClasses.items()
+
     def addFromConfig(self, config: Union[StorageClassConfig, Config, str]) -> None:
         """Add more `StorageClass` definitions from a config file.
 

tests/test_storageClass.py

@@ -143,6 +143,14 @@ def testEquality(self):
         )
         self.assertNotEqual(sc5, sc9)
 
+    def testTypeEquality(self):
+        sc1 = StorageClass("Something", pytype=dict)
+        self.assertTrue(sc1.is_type(dict), repr(sc1))
+        self.assertFalse(sc1.is_type(str), repr(sc1))
+
+        sc2 = StorageClass("TestImage2", "lsst.daf.butler.core.storageClass.StorageClassFactory")
+        self.assertTrue(sc2.is_type(StorageClassFactory), repr(sc2))
+
     def testRegistry(self):
         """Check that storage classes can be created on the fly and stored
         in a registry."""
@@ -164,6 +172,20 @@ def testRegistry(self):
         self.assertNotIn("Temporary3", factory)
         self.assertNotIn({}, factory)
 
+        # Make sure iterators work.
+        keys = set(factory.keys())
+        self.assertIn("Temporary2", keys)
+
+        iterkeys = {k for k in factory}
+        self.assertEqual(keys, iterkeys)
+
+        values = set(factory.values())
+        self.assertIn(sc, values)
+        self.assertEqual(len(factory), len(values))
+
+        external = {k: v for k, v in factory.items()}
+        self.assertIn("Temporary2", external)
+
         # Make sure we can't register a storage class with the same name
         # but different values
         newclass3 = StorageClass("Temporary2", pytype=dict)