Add tests for storageclass conversion on get

python/lsst/daf/butler/tests/_examplePythonTypes.py

@@ -25,10 +25,14 @@
 formats such as FITS or HDF5.
 """
 
-__all__ = ("ListDelegate", "MetricsDelegate", "MetricsExample", "registerMetricsExample")
+__all__ = ("ListDelegate", "MetricsDelegate", "MetricsExample", "registerMetricsExample",
+           "MetricsExampleModel")
 
 
 import copy
+from typing import Optional, Any, Dict, List
+
+from pydantic import BaseModel
 from lsst.daf.butler import StorageClassDelegate, StorageClass
 
 
@@ -234,6 +238,19 @@ def makeFromDict(cls, exportDict):
         return cls(exportDict["summary"], exportDict["output"], data)
 
 
+class MetricsExampleModel(BaseModel):
+    """A variant of `MetricsExample` based on model."""
+
+    summary: Optional[Dict[str, Any]]
+    output: Optional[Dict[str, Any]]
+    data: Optional[List[Any]]
+
+    @classmethod
+    def from_metrics(cls, metrics: MetricsExample) -> "MetricsExampleModel":
+        """Create a model based on an example."""
+        return cls.parse_obj(metrics.exportAsDict())
+
+
 class ListDelegate(StorageClassDelegate):
     """Parameter handler for list parameters"""
 

tests/config/basic/formatters.yaml

@@ -36,3 +36,4 @@ instrument<DummyHSC>:
 StructuredCompositeReadComp: lsst.daf.butler.tests.testFormatters.MetricsExampleFormatter
 StructuredCompositeReadCompNoDisassembly: lsst.daf.butler.tests.testFormatters.MetricsExampleFormatter
 StructuredDataDataTest: lsst.daf.butler.tests.testFormatters.MetricsExampleDataFormatter
+StructuredDataNoComponentsModel: lsst.daf.butler.formatters.json.JsonFormatter

tests/config/basic/storageClasses.yaml

@@ -79,3 +79,11 @@ storageClasses:
       counter: Integer
   StructuredCompositeReadCompNoDisassembly:
     inheritsFrom: StructuredCompositeReadComp
+  StructuredDataNoComponentsModel:
+    # Reading and writing a blob and no components known
+    pytype: lsst.daf.butler.tests.MetricsExampleModel
+    converters:
+      lsst.daf.butler.tests.MetricsExample: lsst.daf.butler.tests.MetricsExampleModel.from_metrics
+      # Add some entries that will fail to import.
+      lsst.daf.butler.bad.type: lsst.daf.butler.tests.MetricsExampleModel.from_metrics
+      lsst.daf.butler.tests.MetricsExampleModel: lsst.daf.butler.bad.function

tests/test_butler.py

@@ -57,6 +57,7 @@ def mock_s3(cls):
 from threading import Thread
 from tempfile import gettempdir
 from lsst.utils import doImport
+from lsst.utils.introspection import get_full_type_name
 from lsst.daf.butler import Butler, Config, ButlerConfig
 from lsst.daf.butler import StorageClassFactory
 from lsst.daf.butler import DatasetType, DatasetRef, DatasetIdGenEnum
@@ -1286,6 +1287,59 @@ def testPruneDatasets(self):
         # Clear out the datasets from registry.
         butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)
 
+    def testPytypeCoercion(self):
+        """Test python type coercion on Butler.get"""
+
+        # Store some data with the normal example storage class.
+        storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
+        datasetTypeName = "test_metric"
+        butler = self.runPutGetTest(storageClass, datasetTypeName)
+
+        dataId = {"instrument": "DummyCamComp", "visit": 423}
+        metric = butler.get(datasetTypeName, dataId=dataId)
+        self.assertEqual(get_full_type_name(metric), "lsst.daf.butler.tests.MetricsExample")
+
+        datasetType_ori = butler.registry.getDatasetType(datasetTypeName)
+        self.assertEqual(datasetType_ori.storageClass.name, "StructuredDataNoComponents")
+
+        # Now need to hack the registry dataset type definition.
+        # There is no API for this.
+        manager = butler.registry._managers.datasets
+        manager._db.update(
+            manager._static.dataset_type,
+            {"name": datasetTypeName},
+            {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataNoComponentsModel"},
+        )
+
+        # Force reset of dataset type cache
+        butler.registry.refresh()
+
+        datasetType_new = butler.registry.getDatasetType(datasetTypeName)
+        self.assertEqual(datasetType_new.name, datasetType_ori.name)
+        self.assertEqual(datasetType_new.storageClass.name, "StructuredDataNoComponentsModel")
+
+        metric_model = butler.get(datasetTypeName, dataId=dataId)
+        self.assertNotEqual(type(metric_model), type(metric))
+        self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel")
+
+        # Put the model and read it back to show that everything now
+        # works as normal.
+        metric_ref = butler.put(metric_model, datasetTypeName, dataId=dataId, visit=424)
+        metric_model_new = butler.get(metric_ref)
+        self.assertEqual(metric_model_new, metric_model)
+
+        # Hack the storage class again to something that will fail on the
+        # get with no conversion class.
+        manager._db.update(
+            manager._static.dataset_type,
+            {"name": datasetTypeName},
+            {datasetTypeName: datasetTypeName, "storage_class": "StructuredDataListYaml"},
+        )
+        butler.registry.refresh()
+
+        with self.assertRaises(ValueError):
+            butler.get(datasetTypeName, dataId=dataId)
+
 
 class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase):
     """InMemoryDatastore specialization of a butler"""