1__all__ = ["TopologicalSorter", "CycleError"]
2
3_NODE_OUT = -1
4_NODE_DONE = -2
5
6
7class _NodeInfo:
8    __slots__ = "node", "npredecessors", "successors"
9
10    def __init__(self, node):
11        # The node this class is augmenting.
12        self.node = node
13
14        # Number of predecessors, generally >= 0. When this value falls to 0,
15        # and is returned by get_ready(), this is set to _NODE_OUT and when the
16        # node is marked done by a call to done(), set to _NODE_DONE.
17        self.npredecessors = 0
18
19        # List of successor nodes. The list can contain duplicated elements as
20        # long as they're all reflected in the successor's npredecessors attribute).
21        self.successors = []
22
23
24class CycleError(ValueError):
25    """Subclass of ValueError raised by TopologicalSorter.prepare if cycles
26    exist in the working graph.
27
28    If multiple cycles exist, only one undefined choice among them will be reported
29    and included in the exception. The detected cycle can be accessed via the second
30    element in the *args* attribute of the exception instance and consists in a list
31    of nodes, such that each node is, in the graph, an immediate predecessor of the
32    next node in the list. In the reported list, the first and the last node will be
33    the same, to make it clear that it is cyclic.
34    """
35
36    pass
37
38
39class TopologicalSorter:
40    """Provides functionality to topologically sort a graph of hashable nodes"""
41
42    def __init__(self, graph=None):
43        self._node2info = {}
44        self._ready_nodes = None
45        self._npassedout = 0
46        self._nfinished = 0
47
48        if graph is not None:
49            for node, predecessors in graph.items():
50                self.add(node, *predecessors)
51
52    def _get_nodeinfo(self, node):
53        if (result := self._node2info.get(node)) is None:
54            self._node2info[node] = result = _NodeInfo(node)
55        return result
56
57    def add(self, node, *predecessors):
58        """Add a new node and its predecessors to the graph.
59
60        Both the *node* and all elements in *predecessors* must be hashable.
61
62        If called multiple times with the same node argument, the set of dependencies
63        will be the union of all dependencies passed in.
64
65        It is possible to add a node with no dependencies (*predecessors* is not provided)
66        as well as provide a dependency twice. If a node that has not been provided before
67        is included among *predecessors* it will be automatically added to the graph with
68        no predecessors of its own.
69
70        Raises ValueError if called after "prepare".
71        """
72        if self._ready_nodes is not None:
73            raise ValueError("Nodes cannot be added after a call to prepare()")
74
75        # Create the node -> predecessor edges
76        nodeinfo = self._get_nodeinfo(node)
77        nodeinfo.npredecessors += len(predecessors)
78
79        # Create the predecessor -> node edges
80        for pred in predecessors:
81            pred_info = self._get_nodeinfo(pred)
82            pred_info.successors.append(node)
83
84    def prepare(self):
85        """Mark the graph as finished and check for cycles in the graph.
86
87        If any cycle is detected, "CycleError" will be raised, but "get_ready" can
88        still be used to obtain as many nodes as possible until cycles block more
89        progress. After a call to this function, the graph cannot be modified and
90        therefore no more nodes can be added using "add".
91        """
92        if self._ready_nodes is not None:
93            raise ValueError("cannot prepare() more than once")
94
95        self._ready_nodes = [
96            i.node for i in self._node2info.values() if i.npredecessors == 0
97        ]
98        # ready_nodes is set before we look for cycles on purpose:
99        # if the user wants to catch the CycleError, that's fine,
100        # they can continue using the instance to grab as many
101        # nodes as possible before cycles block more progress
102        cycle = self._find_cycle()
103        if cycle:
104            raise CycleError(f"nodes are in a cycle", cycle)
105
106    def get_ready(self):
107        """Return a tuple of all the nodes that are ready.
108
109        Initially it returns all nodes with no predecessors; once those are marked
110        as processed by calling "done", further calls will return all new nodes that
111        have all their predecessors already processed. Once no more progress can be made,
112        empty tuples are returned.
113
114        Raises ValueError if called without calling "prepare" previously.
115        """
116        if self._ready_nodes is None:
117            raise ValueError("prepare() must be called first")
118
119        # Get the nodes that are ready and mark them
120        result = tuple(self._ready_nodes)
121        n2i = self._node2info
122        for node in result:
123            n2i[node].npredecessors = _NODE_OUT
124
125        # Clean the list of nodes that are ready and update
126        # the counter of nodes that we have returned.
127        self._ready_nodes.clear()
128        self._npassedout += len(result)
129
130        return result
131
132    def is_active(self):
133        """Return ``True`` if more progress can be made and ``False`` otherwise.
134
135        Progress can be made if cycles do not block the resolution and either there
136        are still nodes ready that haven't yet been returned by "get_ready" or the
137        number of nodes marked "done" is less than the number that have been returned
138        by "get_ready".
139
140        Raises ValueError if called without calling "prepare" previously.
141        """
142        if self._ready_nodes is None:
143            raise ValueError("prepare() must be called first")
144        return self._nfinished < self._npassedout or bool(self._ready_nodes)
145
146    def __bool__(self):
147        return self.is_active()
148
149    def done(self, *nodes):
150        """Marks a set of nodes returned by "get_ready" as processed.
151
152        This method unblocks any successor of each node in *nodes* for being returned
153        in the future by a call to "get_ready".
154
155        Raises :exec:`ValueError` if any node in *nodes* has already been marked as
156        processed by a previous call to this method, if a node was not added to the
157        graph by using "add" or if called without calling "prepare" previously or if
158        node has not yet been returned by "get_ready".
159        """
160
161        if self._ready_nodes is None:
162            raise ValueError("prepare() must be called first")
163
164        n2i = self._node2info
165
166        for node in nodes:
167
168            # Check if we know about this node (it was added previously using add()
169            if (nodeinfo := n2i.get(node)) is None:
170                raise ValueError(f"node {node!r} was not added using add()")
171
172            # If the node has not being returned (marked as ready) previously, inform the user.
173            stat = nodeinfo.npredecessors
174            if stat != _NODE_OUT:
175                if stat >= 0:
176                    raise ValueError(
177                        f"node {node!r} was not passed out (still not ready)"
178                    )
179                elif stat == _NODE_DONE:
180                    raise ValueError(f"node {node!r} was already marked done")
181                else:
182                    assert False, f"node {node!r}: unknown status {stat}"
183
184            # Mark the node as processed
185            nodeinfo.npredecessors = _NODE_DONE
186
187            # Go to all the successors and reduce the number of predecessors, collecting all the ones
188            # that are ready to be returned in the next get_ready() call.
189            for successor in nodeinfo.successors:
190                successor_info = n2i[successor]
191                successor_info.npredecessors -= 1
192                if successor_info.npredecessors == 0:
193                    self._ready_nodes.append(successor)
194            self._nfinished += 1
195
196    def _find_cycle(self):
197        n2i = self._node2info
198        stack = []
199        itstack = []
200        seen = set()
201        node2stacki = {}
202
203        for node in n2i:
204            if node in seen:
205                continue
206
207            while True:
208                if node in seen:
209                    # If we have seen already the node and is in the
210                    # current stack we have found a cycle.
211                    if node in node2stacki:
212                        return stack[node2stacki[node] :] + [node]
213                    # else go on to get next successor
214                else:
215                    seen.add(node)
216                    itstack.append(iter(n2i[node].successors).__next__)
217                    node2stacki[node] = len(stack)
218                    stack.append(node)
219
220                # Backtrack to the topmost stack entry with
221                # at least another successor.
222                while stack:
223                    try:
224                        node = itstack[-1]()
225                        break
226                    except StopIteration:
227                        del node2stacki[stack.pop()]
228                        itstack.pop()
229                else:
230                    break
231        return None
232
233    def static_order(self):
234        """Returns an iterable of nodes in a topological order.
235
236        The particular order that is returned may depend on the specific
237        order in which the items were inserted in the graph.
238
239        Using this method does not require to call "prepare" or "done". If any
240        cycle is detected, :exc:`CycleError` will be raised.
241        """
242        self.prepare()
243        while self.is_active():
244            node_group = self.get_ready()
245            yield from node_group
246            self.done(*node_group)
247