Move composites logic shared by pickle and JSON to composites.py

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

@@ -22,6 +22,8 @@
 
 """Support for reading and writing composite objects."""
 
+import collections
+
 
 class DatasetComponent:
 
@@ -124,6 +126,35 @@ def genericAssembler(storageClass, components, pytype=None):
     return obj
 
 
+def validComponents(composite, storageClass):
+    """Extract requested components from a composite.
+
+    Parameters
+    ----------
+    composite : `object`
+        Composite from which to extract components.
+    storageClass : `StorageClass`
+        `StorageClass` associated with this composite object.
+        If None, it is assumed this is not to be treated as a composite.
+
+    Returns
+    -------
+    comps : `dict`
+        Non-None components extracted from the composite, indexed by the component
+        name as derived from the `StorageClass`.
+    """
+    components = {}
+    if storageClass is not None and storageClass.components:
+        for c in storageClass.components:
+            if isinstance(composite, collections.Mapping):
+                comp = composite[c]
+            else:
+                comp = genericGetter(composite, c)
+            if comp is not None:
+                components[c] = comp
+    return components
+
+
 def hasComponent(composite, componentName):
     """Determine if it seems likely that the composite has the component.
 

python/lsst/daf/butler/formatters/jsonFormatter.py

@@ -25,7 +25,7 @@
 import json
 import collections
 
-from lsst.daf.butler.core.composites import genericAssembler, hasComponent
+from lsst.daf.butler.core.composites import genericAssembler, validComponents
 from lsst.daf.butler.core.formatter import Formatter
 
 
@@ -156,15 +156,8 @@ def write(self, inMemoryDataset, fileDescriptor):
                 inMemoryDataset = inMemoryDataset._asdict()
             json.dump(inMemoryDataset, fd)
 
-        # The reference components come from the StorageClass when determining
-        # whether this object can be a composite.
-        sc = fileDescriptor.storageClass
-        comps = []
-        if sc is not None and sc.components:
-            for c in sc.components:
-                if (hasComponent(inMemoryDataset, c) is not None or
-                        (isinstance(inMemoryDataset, collections.Mapping) and c in inMemoryDataset)):
-                    comps.append(c)
+        # Get the list of valid components so we can build URIs
+        components = validComponents(inMemoryDataset, fileDescriptor.storageClass)
 
         return (fileDescriptor.location.uri,
-                {c: fileDescriptor.location.componentUri(c) for c in comps})
+                {c: fileDescriptor.location.componentUri(c) for c in components})

python/lsst/daf/butler/formatters/pickleFormatter.py

@@ -26,7 +26,7 @@
 import pickle
 import collections
 
-from lsst.daf.butler.core.composites import genericGetter, hasComponent
+from lsst.daf.butler.core.composites import genericGetter, validComponents
 from lsst.daf.butler.core.formatter import Formatter
 
 
@@ -151,15 +151,8 @@ def write(self, inMemoryDataset, fileDescriptor):
         with open(self._getPath(fileDescriptor), "wb") as fd:
             pickle.dump(inMemoryDataset, fd, protocol=-1)
 
-        # The reference components come from the StorageClass when determining
-        # whether this object can be a composite.
-        sc = fileDescriptor.storageClass
-        comps = []
-        if sc is not None and sc.components:
-            for c in sc.components:
-                if (hasComponent(inMemoryDataset, c) is not None or
-                        (isinstance(inMemoryDataset, collections.Mapping) and c in inMemoryDataset)):
-                    comps.append(c)
+        # Get the list of valid components so we can build URIs
+        components = validComponents(inMemoryDataset, fileDescriptor.storageClass)
 
         return (fileDescriptor.location.uri,
-                {c: fileDescriptor.location.componentUri(c) for c in comps})
+                {c: fileDescriptor.location.componentUri(c) for c in components})