1414
1515from __future__ import annotations
1616
17+ import heapq
1718from collections .abc import Generator , Iterator
1819
1920from . import _typings as t
2526 class CanHashAndCompareLT (typing .Protocol ):
2627 def __hash__ (self ) -> int : ...
2728
28- def __lt__ (self , other : typing .Self , / ) -> bool : ...
29+ def __lt__ (self , other : typing .Any , / ) -> bool : ...
2930
3031 class CanHashAndCompareGT (typing .Protocol ):
3132 def __hash__ (self ) -> int : ...
3233
33- def __gt__ (self , other : typing .Self , / ) -> bool : ...
34+ def __gt__ (self , other : typing .Any , / ) -> bool : ...
3435
3536else :
3637
@@ -62,7 +63,7 @@ def cycle(self) -> list[T]:
6263 return self .args [0 ]
6364
6465
65- class NodeData [T ]:
66+ class NodeData [T : CanHashAndCompare ]:
6667 __slots__ = ("dependants" , "ndependencies" , "node" )
6768
6869 def __init__ (self , node : T ) -> None :
@@ -74,6 +75,9 @@ def __init_subclass__(cls) -> t.Never:
7475 msg = "Don't subclass this"
7576 raise RuntimeError (msg )
7677
78+ def __lt__ (self , other : t .Self ) -> bool : # falback for tuple sorting
79+ return True
80+
7781 __final__ = True
7882
7983
@@ -174,14 +178,17 @@ def __iter__(self) -> Generator[T, None, None]:
174178 return self .__iter ()
175179
176180 def __iter (self ) -> Generator [T , None , None ]:
177- while ready := [
178- i .node for i in self ._nodemap .values () if not i .ndependencies
179- ]:
180- next_node = min (ready )
181- self ._nodemap [next_node ].ndependencies = - 1
181+ m = self ._nodemap
182+ ready = [(n , i ) for n , i in m .items () if not i .ndependencies ]
183+ heapq .heapify (ready )
184+ while ready :
185+ next_node , info = heapq .heappop (ready )
186+ info .ndependencies = - 1
182187
183188 yield next_node
184189
185- for dep in self . _nodemap [ next_node ] .dependants :
190+ for dep in info .dependants :
186191 dep_info = self ._nodemap [dep ]
187192 dep_info .ndependencies -= 1
193+ if not dep_info .ndependencies :
194+ heapq .heappush (ready , (dep , dep_info ))
0 commit comments