Add some tests for DatasetType parentStorageClass

This required changes to how pickling and deep copy worked.
Also to simplify pickling the parentStorageClass is now
also a positional argument.

python/lsst/daf/butler/core/datasets/type.py

@@ -117,7 +117,7 @@ def nameWithComponent(datasetTypeName: str, componentName: str) -> str:
 
     def __init__(self, name: str, dimensions: Union[DimensionGraph, Iterable[Dimension]],
                  storageClass: Union[StorageClass, str],
-                 *, parentStorageClass: Optional[Union[StorageClass, str]] = None,
+                 parentStorageClass: Optional[Union[StorageClass, str]] = None, *,
                  universe: Optional[DimensionUniverse] = None):
         if self.VALID_NAME_REGEX.match(name) is None:
             raise ValueError(f"DatasetType name '{name}' is invalid.")
@@ -128,7 +128,9 @@ def __init__(self, name: str, dimensions: Union[DimensionGraph, Iterable[Dimensi
                                  "a universe must be provided.")
             dimensions = universe.extract(dimensions)
         self._dimensions = dimensions
-        assert isinstance(storageClass, (StorageClass, str))
+        if not isinstance(storageClass, (StorageClass, str)):
+            raise ValueError("StorageClass argument must be StorageClass or str. "
+                             f"Got {storageClass}")
         self._storageClass: Optional[StorageClass]
         if isinstance(storageClass, StorageClass):
             self._storageClass = storageClass
@@ -140,6 +142,10 @@ def __init__(self, name: str, dimensions: Union[DimensionGraph, Iterable[Dimensi
         self._parentStorageClass: Optional[StorageClass] = None
         self._parentStorageClassName: Optional[str] = None
         if parentStorageClass is not None:
+            if not isinstance(storageClass, (StorageClass, str)):
+                raise ValueError("Parent StorageClass argument must be StorageClass or str. "
+                                 f"Got {parentStorageClass}")
+
             # Only allowed for a component dataset type
             _, componentName = self.splitDatasetTypeName(self._name)
             if componentName is None:
@@ -155,13 +161,16 @@ def __init__(self, name: str, dimensions: Union[DimensionGraph, Iterable[Dimensi
         # a component and is not specified when we don't
         _, componentName = self.splitDatasetTypeName(self._name)
         if parentStorageClass is None and componentName is not None:
-            raise RuntimeError(f"Component dataset type '{self._name}' constructed without parent"
-                               " storage class")
+            raise ValueError(f"Component dataset type '{self._name}' constructed without parent"
+                             " storage class")
         if parentStorageClass is not None and componentName is None:
-            raise RuntimeError(f"Parent storage class specified by {self._name} is not a composite")
+            raise ValueError(f"Parent storage class specified by {self._name} is not a composite")
 
     def __repr__(self) -> str:
-        return "DatasetType({}, {}, {})".format(self.name, self.dimensions, self._storageClassName)
+        parent = ""
+        if self._parentStorageClassName:
+            parent = f", parentStorageClass={self._parentStorageClassName}"
+        return f"DatasetType({self.name}, {self.dimensions}, {self._storageClassName}{parent})"
 
     def __eq__(self, other: Any) -> bool:
         if not isinstance(other, type(self)):
@@ -171,17 +180,24 @@ def __eq__(self, other: Any) -> bool:
         if self._dimensions != other._dimensions:
             return False
         if self._storageClass is not None and other._storageClass is not None:
-            return self._storageClass == other._storageClass
+            if self._storageClass != other._storageClass:
+                return False
+        else:
+            if self._storageClassName != other._storageClassName:
+                return False
+        if self._parentStorageClass is not None and other._parentStorageClass is not None:
+            return self._parentStorageClass == other._parentStorageClass
         else:
-            return self._storageClassName == other._storageClassName
+            return self._parentStorageClassName == other._parentStorageClassName
 
     def __hash__(self) -> int:
         """Hash DatasetType instance.
 
         This only uses StorageClass name which is it consistent with the
         implementation of StorageClass hash method.
         """
-        return hash((self._name, self._dimensions, self._storageClassName))
+        return hash((self._name, self._dimensions, self._storageClassName,
+                     self._parentStorageClassName))
 
     @property
     def name(self) -> str:
@@ -382,13 +398,14 @@ def _lookupNames(self) -> Tuple[LookupKey, ...]:
 
         return lookups + self.storageClass._lookupNames()
 
-    def __reduce__(self) -> Tuple[Type[DatasetType], Tuple[str, DimensionGraph, str]]:
+    def __reduce__(self) -> Tuple[Type[DatasetType], Tuple[str, DimensionGraph, str, Optional[str]]]:
         """Support pickling.
 
         StorageClass instances can not normally be pickled, so we pickle
         StorageClass name instead of instance.
         """
-        return (DatasetType, (self.name, self.dimensions, self._storageClassName))
+        return (DatasetType, (self.name, self.dimensions, self._storageClassName,
+                              self._parentStorageClassName))
 
     def __deepcopy__(self, memo: Any) -> DatasetType:
         """Support for deep copy method.
@@ -401,4 +418,6 @@ def __deepcopy__(self, memo: Any) -> DatasetType:
         """
         return DatasetType(name=deepcopy(self.name, memo),
                            dimensions=deepcopy(self.dimensions, memo),
-                           storageClass=deepcopy(self._storageClass or self._storageClassName, memo))
+                           storageClass=deepcopy(self._storageClass or self._storageClassName, memo),
+                           parentStorageClass=deepcopy(self._parentStorageClass
+                                                            or self._parentStorageClassName, memo))

tests/test_datasets.py

@@ -21,6 +21,7 @@
 
 import unittest
 import pickle
+import copy
 
 from lsst.daf.butler import (
     DataCoordinate,
@@ -55,6 +56,13 @@ def testConstructor(self):
         self.assertEqual(datasetType.storageClass, storageClass)
         self.assertEqual(datasetType.dimensions, dimensions)
 
+        with self.assertRaises(ValueError, msg="Construct component without parent storage class"):
+            DatasetType(DatasetType.nameWithComponent(datasetTypeName, "comp"),
+                        dimensions, storageClass)
+        with self.assertRaises(ValueError, msg="Construct non-component with parent storage class"):
+            DatasetType(datasetTypeName,
+                        dimensions, storageClass, parentStorageClass="NotAllowed")
+
     def testConstructor2(self):
         """Test construction from StorageClass name.
         """
@@ -162,6 +170,23 @@ def testHashability(self):
         self.assertNotEqual(hash(DatasetType("a", dimensions, "test_c")),
                             hash(DatasetType("a", dimensions, "test_d")))
 
+    def testDeepCopy(self):
+        """Test that we can copy a dataset type."""
+        storageClass = StorageClass("test_copy")
+        datasetTypeName = "test"
+        dimensions = self.universe.extract(("instrument", "visit"))
+        datasetType = DatasetType(datasetTypeName, dimensions, storageClass)
+        dcopy = copy.deepcopy(datasetType)
+        self.assertEqual(dcopy, datasetType)
+
+        # And again with a composite
+        componentStorageClass = StorageClass("copy_component")
+        componentDatasetType = DatasetType(DatasetType.nameWithComponent(datasetTypeName, "comp"),
+                                           dimensions, componentStorageClass,
+                                           parentStorageClass=storageClass)
+        dcopy = copy.deepcopy(componentDatasetType)
+        self.assertEqual(dcopy, componentDatasetType)
+
     def testPickle(self):
         """Test pickle support.
         """
@@ -176,6 +201,44 @@ def testPickle(self):
         self.assertEqual(datasetType.name, datasetTypeOut.name)
         self.assertEqual(datasetType.dimensions.names, datasetTypeOut.dimensions.names)
         self.assertEqual(datasetType.storageClass, datasetTypeOut.storageClass)
+        self.assertIsNone(datasetTypeOut.parentStorageClass)
+
+        # And again with a composite
+        componentStorageClass = StorageClass("pickle_component")
+        StorageClassFactory().registerStorageClass(componentStorageClass)
+        componentDatasetType = DatasetType(DatasetType.nameWithComponent(datasetTypeName, "comp"),
+                                           dimensions, componentStorageClass,
+                                           parentStorageClass=storageClass)
+        datasetTypeOut = pickle.loads(pickle.dumps(componentDatasetType))
+        self.assertIsInstance(datasetTypeOut, DatasetType)
+        self.assertEqual(componentDatasetType.name, datasetTypeOut.name)
+        self.assertEqual(componentDatasetType.dimensions.names, datasetTypeOut.dimensions.names)
+        self.assertEqual(componentDatasetType.storageClass, datasetTypeOut.storageClass)
+        self.assertEqual(componentDatasetType.parentStorageClass, datasetTypeOut.parentStorageClass)
+        self.assertEqual(datasetTypeOut.parentStorageClass.name,
+                         storageClass.name)
+        self.assertEqual(datasetTypeOut, componentDatasetType)
+
+        # Now with a string and not a real storage class to test that
+        # pickling doesn't force the StorageClass to be resolved
+        componentDatasetType = DatasetType(DatasetType.nameWithComponent(datasetTypeName, "comp"),
+                                           dimensions, "StrangeComponent",
+                                           parentStorageClass="UnknownParent")
+        datasetTypeOut = pickle.loads(pickle.dumps(componentDatasetType))
+        self.assertEqual(datasetTypeOut, componentDatasetType)
+        self.assertEqual(datasetTypeOut._parentStorageClassName,
+                         componentDatasetType._parentStorageClassName)
+
+        # Now with a storage class that is created by the factory
+        factoryStorageClassClass = StorageClassFactory.makeNewStorageClass("ParentClass")
+        factoryComponentStorageClassClass = StorageClassFactory.makeNewStorageClass("ComponentClass")
+        componentDatasetType = DatasetType(DatasetType.nameWithComponent(datasetTypeName, "comp"),
+                                           dimensions, factoryComponentStorageClassClass(),
+                                           parentStorageClass=factoryStorageClassClass())
+        datasetTypeOut = pickle.loads(pickle.dumps(componentDatasetType))
+        self.assertEqual(datasetTypeOut, componentDatasetType)
+        self.assertEqual(datasetTypeOut._parentStorageClassName,
+                         componentDatasetType._parentStorageClassName)
 
     def test_composites(self):
         """Test components within composite DatasetTypes."""
@@ -204,6 +267,10 @@ def test_composites(self):
         self.assertEqual(datasetTypeComposite.nameAndComponent(), ("composite", None))
         self.assertEqual(datasetTypeComponentA.nameAndComponent(), ("composite", "compA"))
 
+        self.assertEqual(datasetTypeComponentA.parentStorageClass, storageClass)
+        self.assertEqual(datasetTypeComponentB.parentStorageClass, storageClass)
+        self.assertIsNone(datasetTypeComposite.parentStorageClass)
+
 
 class DatasetRefTestCase(unittest.TestCase):
     """Test for DatasetRef.