Implements a routine for topologically sorting a DAG with dynamic keying function

Author: kaushikcfd

pytools/graph.py

@@ -33,9 +33,11 @@
 .. autofunction:: compute_sccs
 .. autoclass:: CycleError
 .. autofunction:: compute_topological_order
+.. autofunction:: compute_topological_order_with_dynamic_key
 .. autofunction:: compute_transitive_closure
 .. autofunction:: contains_cycle
 .. autofunction:: compute_induced_subgraph
+.. autoclass:: TopologicalOrderState
 
 Type Variables Used
 -------------------
@@ -46,7 +48,8 @@
 """
 
 from typing import (TypeVar, Mapping, Iterable, List, Optional, Any, Callable,
-                    Set, MutableSet, Dict, Iterator, Tuple)
+                    Set, MutableSet, Dict, Iterator, Tuple, Generic)
+from dataclasses import dataclass
 
 
 T = TypeVar("T")
@@ -207,18 +210,44 @@ def __lt__(self, other):
         return self.key < other.key
 
 
-def compute_topological_order(graph: Mapping[T, Iterable[T]],
-                              key: Optional[Callable[[T], Any]] = None) -> List[T]:
-    """Compute a topological order of nodes in a directed graph.
+@dataclass(frozen=True)
+class TopologicalOrderState(Generic[T]):
+    """
+    .. attribute:: scheduled_nodes
+
+        A :class:`list` of nodes that have been scheduled.
+
+    .. warning::
+
+        - Mutable updates to :attr:`scheduled_nodes`
+          results in an undefined behavior.
+    """
+    scheduled_nodes: List[T]
+
+
+def compute_topological_order_with_dynamic_key(
+        graph: Mapping[T, Iterable[T]],
+        trigger_key_update: Callable[[TopologicalOrderState[T]], bool],
+        get_key: Callable[[TopologicalOrderState[T]], Callable[[T], Any]]
+) -> List[T]:
+    """
+    Computes a topological order of nodes in a directed graph with support for
+    a dynamic keying function.
 
     :arg graph: A :class:`collections.abc.Mapping` representing a directed
         graph. The dictionary contains one key representing each node in the
         graph, and this key maps to a :class:`collections.abc.Iterable` of its
         successor nodes.
 
-    :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 trigger_key_update: A function called after scheduling a node in
+        *graph* that takes in an instance of :class:`TopologicalOrderState`
+        corresponding to the scheduling state at that point and returns whether
+        the comparison keys corresponding to the nodes be updated.
+
+    :arg get_key: A callable  called when *trigger_key_update*
+        returns *True*. Takes in an instance of :class:`TopologicalOrderState`
+        and returns another callable that accepts node as an argument and returns the
+        comparison key corresponding to the node.
 
     :returns: A :class:`list` representing a valid topological ordering of the
         nodes in the directed graph.
@@ -228,10 +257,8 @@ def compute_topological_order(graph: Mapping[T, Iterable[T]],
         * Requires the keys of the mapping *graph* to be hashable.
         * Implements `Kahn's algorithm <https://w.wiki/YDy>`__.
 
-    .. versionadded:: 2020.2
+    .. versionadded:: 2022.2
     """
-    # all nodes have the same keys when not provided
-    keyfunc = key if key is not None else (lambda x: 0)
 
     from heapq import heapify, heappop, heappush
 
@@ -248,6 +275,8 @@ def compute_topological_order(graph: Mapping[T, Iterable[T]],
 
     # }}}
 
+    keyfunc = get_key(TopologicalOrderState(scheduled_nodes=[]))
+
     total_num_nodes = len(nodes_to_num_predecessors)
 
     # heap: list of instances of HeapEntry(n) where 'n' is a node in
@@ -263,6 +292,14 @@ def compute_topological_order(graph: Mapping[T, Iterable[T]],
         node_to_be_scheduled = heappop(heap).node
         order.append(node_to_be_scheduled)
 
+        state = TopologicalOrderState(scheduled_nodes=order)
+
+        if trigger_key_update(state):
+            keyfunc = get_key(state)
+            heap = [HeapEntry(entry.node, keyfunc(entry.node))
+                    for entry in heap]
+            heapify(heap)
+
         # discard 'node_to_be_scheduled' from the predecessors of its
         # successors since it's been scheduled
         for child in graph.get(node_to_be_scheduled, ()):
@@ -277,6 +314,38 @@ def compute_topological_order(graph: Mapping[T, Iterable[T]],
 
     return order
 
+
+def compute_topological_order(graph: Mapping[T, Iterable[T]],
+                              key: Optional[Callable[[T], Any]] = None) -> List[T]:
+    """Compute a topological order of nodes in a directed graph.
+
+    :arg graph: A :class:`collections.abc.Mapping` representing a directed
+        graph. The dictionary contains one key representing each node in the
+        graph, and this key maps to a :class:`collections.abc.Iterable` of its
+        successor nodes.
+
+    :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*.
+
+    :returns: A :class:`list` representing a valid topological ordering of the
+        nodes in the directed graph.
+
+    .. note::
+
+        * Requires the keys of the mapping *graph* to be hashable.
+        * Implements `Kahn's algorithm <https://w.wiki/YDy>`__.
+
+    .. versionadded:: 2020.2
+    """
+    # all nodes have the same keys when not provided
+    keyfunc = key if key is not None else (lambda x: 0)
+
+    return compute_topological_order_with_dynamic_key(
+        graph,
+        trigger_key_update=lambda _: False,
+        get_key=lambda _: keyfunc)
+
 # }}}