Better cycle detection

pytools/graph.py

@@ -41,6 +41,7 @@
 .. autoexception:: CycleError
 .. autofunction:: compute_topological_order
 .. autofunction:: compute_transitive_closure
+.. autofunction:: find_cycles
 .. autofunction:: contains_cycle
 .. autofunction:: compute_induced_subgraph
 .. autofunction:: validate_graph
@@ -240,6 +241,42 @@ def __init__(self, node: NodeT) -> None:
         self.node = node
 
 
+def find_cycles(graph: GraphT) -> List[List[NodeT]]:
+    """
+    Find all cycles in *graph* using DFS.
+
+    :returns: A :class:`list` in which each element represents another :class:`list`
+        of nodes that form a cycle.
+    """
+    def dfs(node: NodeT, path: List[NodeT]) -> List[NodeT]:
+        # Cycle detected
+        if visited[node] == 1:
+            return path
+
+        # Visit this node, explore its children
+        visited[node] = 1
+        path.append(node)
+        for child in graph[node]:
+            if visited[child] != 2 and dfs(child, path):
+                return path
+
+        # Done visiting node
+        visited[node] = 2
+        return []
+
+    visited = {node: 0 for node in graph.keys()}
+
+    res = []
+
+    for node in graph:
+        if not visited[node]:
+            cycle = dfs(node, [])
+            if cycle:
+                res.append(cycle)
+
+    return res
+
+
 class HeapEntry:
     """
     Helper class to compare associated keys while comparing the elements in
@@ -257,14 +294,16 @@ def __lt__(self, other: "HeapEntry") -> bool:
 
 
 def compute_topological_order(graph: GraphT,
-                              key: Optional[Callable[[NodeT], Any]] = None) \
-                              -> List[NodeT]:
+                              key: Optional[Callable[[NodeT], Any]] = None,
+                              verbose_cycle: bool = True) -> List[NodeT]:
     """Compute a topological order of nodes in a directed graph.
 
     :arg key: A custom key function may be supplied to determine the order in
         break-even cases. Expects a function of one argument that is used to
         extract a comparison key from each node of the *graph*.
 
+    :arg verbose_cycle: Verbose reporting in case *graph* contains a cycle.
+
     :returns: A :class:`list` representing a valid topological ordering of the
         nodes in the directed graph.
 
@@ -316,9 +355,17 @@ def compute_topological_order(graph: GraphT,
                 heappush(heap, HeapEntry(child, keyfunc(child)))
 
     if len(order) != total_num_nodes:
-        # any node which has a predecessor left is a part of a cycle
-        raise CycleError(next(iter(n for n, num_preds in
-            nodes_to_num_predecessors.items() if num_preds != 0)))
+        # There is a cycle in the graph
+        if not verbose_cycle:
+            raise CycleError(None)
+
+        try:
+            cycles = find_cycles(graph)
+        except KeyError:
+            # Graph is invalid
+            raise CycleError(None)
+        else:
+            raise CycleError(cycles[0][0])
 
     return order
 
@@ -371,7 +418,7 @@ def contains_cycle(graph: GraphT) -> bool:
     """
 
     try:
-        compute_topological_order(graph)
+        compute_topological_order(graph, verbose_cycle=False)
         return False
     except CycleError:
         return True

test/test_graph_tools.py

@@ -395,6 +395,34 @@ def test_is_connected():
     assert is_connected({})
 
 
+def test_find_cycles():
+    from pytools.graph import compute_topological_order, CycleError, find_cycles
+
+    # Non-Self Loop
+    graph = {1: {}, 5: {1, 8}, 8: {5}}
+    assert find_cycles(graph) == [[5, 8]]
+    with pytest.raises(CycleError, match="5|8"):
+        compute_topological_order(graph)
+
+    # Self-Loop
+    graph = {1: {1}, 5: {8}, 8: {}}
+    assert find_cycles(graph) == [[1]]
+    with pytest.raises(CycleError, match="1"):
+        compute_topological_order(graph)
+
+    # Invalid graph with loop
+    graph = {1: {42}, 5: {8}, 8: {5}}
+    # Can't run find_cycles on this graph since it is invalid
+    with pytest.raises(CycleError, match="None"):
+        compute_topological_order(graph)
+
+    # Multiple loops
+    graph = {1: {1}, 5: {8}, 8: {5}}
+    assert find_cycles(graph) == [[1], [5, 8]]
+    with pytest.raises(CycleError, match="1"):
+        compute_topological_order(graph)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])