DM-14544: Add pickle support to Butler

python/lsst/daf/butler/butler.py

@@ -167,6 +167,11 @@ def __init__(self, config=None, collection=None, run=None):
             if self.run is None:
                 self.run = self.registry.makeRun(runCollection)
 
+    def __reduce__(self):
+        """Support pickling.
+        """
+        return (Butler, (self.config, ))
+
     def put(self, obj, datasetType, dataId, producer=None):
         """Store and register a dataset.
 

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

@@ -218,6 +218,11 @@ def __call__(cls):  # noqa N805
         return cls._instances[cls]
 
 
+# Helper Node in a TopologicalSet
+# Unfortunately can't be a class member because that breaks pickling
+TopologicalSetNode = namedtuple('TopologicalSetNode', ['element', 'sourceElements'])
+
+
 class TopologicalSet:
     """A collection that behaves like a builtin `set`, but where
     elements can be interconnected (like a graph).
@@ -230,10 +235,8 @@ class TopologicalSet:
     elements : `iterable`
         Any iterable with elements to insert.
     """
-    Node = namedtuple('Node', ['element', 'sourceElements'])
-
     def __init__(self, elements):
-        self._nodes = {e: TopologicalSet.Node(e, set()) for e in elements}
+        self._nodes = {e: TopologicalSetNode(e, set()) for e in elements}
         self._ordering = None
 
     def __contains__(self, element):
@@ -242,6 +245,9 @@ def __contains__(self, element):
     def __len__(self):
         return len(self._nodes)
 
+    def __eq__(self, other):
+        return self._nodes == other._nodes
+
     def connect(self, sourceElement, targetElement):
         """Connect two elements in the set.
 

tests/test_butler.py

@@ -25,6 +25,7 @@
 import os
 import unittest
 from tempfile import TemporaryDirectory
+import pickle
 
 import lsst.utils.tests
 
@@ -159,6 +160,14 @@ def testMakeRepo(self):
         self.assertIn("datastore.formatters", full)
         self.assertNotIn("datastore.formatters", limited)
 
+    def testPickle(self):
+        """Test pickle support.
+        """
+        butler = Butler(self.configFile)
+        butlerOut = pickle.loads(pickle.dumps(butler))
+        self.assertIsInstance(butlerOut, Butler)
+        self.assertEqual(butlerOut.config, butler.config)
+
 
 class MemoryTester(lsst.utils.tests.MemoryTestCase):
     pass

tests/test_utils.py

@@ -21,6 +21,7 @@
 
 import unittest
 import inspect
+import pickle
 from itertools import permutations
 from random import shuffle
 
@@ -181,6 +182,17 @@ def testTopologicalOrdering(self):
             topologicalSet2.connect(i, i+1)
         self.assertEqual(list(topologicalSet2), elements)
 
+    def testPickle(self):
+        """Should be possible to pickle.
+        """
+        elements = ['a', 'd', 'f']
+        topologicalSet = TopologicalSet(elements)
+        # Adding connections should work
+        topologicalSet.connect('a', 'd')
+        topologicalSet.connect('a', 'f')
+        out = pickle.loads(pickle.dumps(topologicalSet))
+        self.assertEqual(out, topologicalSet)
+
 
 class TestButlerUtils(lsst.utils.tests.TestCase):
     """Tests of the simple utilities."""