Merge pull request #29 from lsst/tickets/DM-14012

DM-14012: Add helper container class supporting topologically sorted iteration over elements

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

@@ -19,6 +19,8 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+from collections import namedtuple
+
 
 def iterable(a):
     """Make input iterable.
@@ -203,3 +205,93 @@ def __call__(cls):  # noqa N805
         if cls not in cls._instances:
             cls._instances[cls] = super(Singleton, cls).__call__()
         return cls._instances[cls]
+
+
+class TopologicalSet:
+    """A collection that behaves like a builtin `set`, but where
+    elements can be interconnected (like a graph).
+
+    Iteration over this collection visits its elements in topologically
+    sorted order.
+
+    Parameters
+    ----------
+    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._ordering = None
+
+    def __contains__(self, element):
+        return element in self._nodes
+
+    def __len__(self):
+        return len(self._nodes)
+
+    def connect(self, sourceElement, targetElement):
+        """Connect two elements in the set.
+
+        The connection is directed from `sourceElement` to `targetElement` and
+        is distinct from its inverse.
+        Both elements must already be present in the set.
+
+        sourceElement : `object`
+            The source element.
+        targetElement : `object`
+            The target element.
+
+        Raises
+        ------
+        KeyError
+            When either element is not already in the set.
+        ValueError
+            If a self connections between elements would be created.
+        """
+        if sourceElement == targetElement:
+            raise ValueError('Cannot connect {} to itself'.format(sourceElement))
+        for element in (sourceElement, targetElement):
+            if element not in self._nodes:
+                raise KeyError('{} not in set'.format(element))
+        targetNode = self._nodes[targetElement]
+        targetNode.sourceElements.add(sourceElement)
+        # Adding a connection invalidates previous ordering
+        self._ordering = None
+
+    def __iter__(self):
+        """Iterate over elements in topologically sorted order.
+
+        Raises
+        ------
+        ValueError
+            If a cycle is found and hence no topological order exists.
+        """
+        if self._ordering is None:
+            self.ordering = self._topologicalOrdering()
+        yield from self.ordering
+
+    def _topologicalOrdering(self):
+        """Generate a topological ordering by doing a basic
+        depth-first-search.
+        """
+        seen = set()
+        finished = set()
+        order = []
+
+        def visit(node):
+            if node.element in finished:
+                return
+            if node.element in seen:
+                raise ValueError("Cycle detected")
+            seen.add(node.element)
+            for sourceElement in node.sourceElements:
+                visit(self._nodes[sourceElement])
+            finished.add(node.element)
+            seen.remove(node.element)
+            order.append(node.element)
+
+        for node in self._nodes.values():
+            visit(node)
+        return order

tests/test_utils.py

@@ -21,10 +21,12 @@
 
 import unittest
 import inspect
+from itertools import permutations
+from random import shuffle
 
 import lsst.utils.tests
 
-from lsst.daf.butler.core.utils import iterable, doImport, getFullTypeName, Singleton
+from lsst.daf.butler.core.utils import iterable, doImport, getFullTypeName, Singleton, TopologicalSet
 from lsst.daf.butler.core.formatter import Formatter
 from lsst.daf.butler import StorageClass
 
@@ -105,6 +107,69 @@ def testSingleton(self):
             SingletonTestCase.IsBadSingleton(52)
 
 
+class TopologicalSetTestCase(lsst.utils.tests.TestCase):
+    """Tests of the TopologicalSet class"""
+
+    def testConstructor(self):
+        elements = [1, 5, 7, 3, 10, 11]
+        topologicalSet = TopologicalSet(elements)
+        for e in elements:
+            self.assertIn(e, topologicalSet)
+        self.assertEqual(len(elements), len(topologicalSet))
+
+    def testConnect(self):
+        elements = ['a', 'd', 'f']
+        topologicalSet = TopologicalSet(elements)
+        # Adding connections should work
+        topologicalSet.connect('a', 'd')
+        topologicalSet.connect('a', 'f')
+        # Even when adding cycles (unfortunately)
+        topologicalSet.connect('f', 'a')
+        # Adding a connection from or to a non existing element should also fail
+        with self.assertRaises(KeyError):
+            topologicalSet.connect('a', 'c')
+        with self.assertRaises(KeyError):
+            topologicalSet.connect('c', 'a')
+        with self.assertRaises(KeyError):
+            topologicalSet.connect('c', 'g')
+        with self.assertRaises(ValueError):
+            topologicalSet.connect('c', 'c')
+
+    def testTopologicalOrdering(self):
+        """Iterating over a TopologicalSet should visit the elements
+        in the set in topologically sorted order.
+        """
+        # First check a basic topological ordering
+        elements = ['shoes', 'belt', 'trousers']
+        for p in permutations(elements):
+            topologicalSet = TopologicalSet(p)
+            topologicalSet.connect('belt', 'shoes')
+            topologicalSet.connect('trousers', 'shoes')
+            # Check valid orderings
+            self.assertIn(list(topologicalSet), [['belt', 'trousers', 'shoes'],
+                                                 ['trousers', 'belt', 'shoes']])
+            # Check invalid orderings (probably redundant)
+            self.assertNotIn(list(topologicalSet), [['shoes', 'belt', 'trousers'],
+                                                    ['shoes', 'trousers', 'belt']])
+        # Adding a cycle should cause iteration to fail
+        topologicalSet.connect('shoes', 'belt')
+        with self.assertRaises(ValueError):
+            ignore = list(topologicalSet)  # noqa F841
+        # Now check for a larger number of elements.
+        # Here we can't possibly test all possible valid topological orderings,
+        # so instead we connect all elements.
+        # Thus the topological sort should be equivalent to a regular sort
+        # (but less efficient).
+        N = 100
+        elements = list(range(N))
+        unorderedElements = elements.copy()
+        shuffle(unorderedElements)
+        topologicalSet2 = TopologicalSet(unorderedElements)
+        for i in range(N-1):
+            topologicalSet2.connect(i, i+1)
+        self.assertEqual(list(topologicalSet2), elements)
+
+
 class TestButlerUtils(lsst.utils.tests.TestCase):
     """Tests of the simple utilities."""