11from __future__ import annotations
22
3+ import sys
34from collections import defaultdict
4- from typing import Any
5- from typing import DefaultDict
6- from typing import Generic
7- from typing import Iterable
8- from typing import Iterator
9- from typing import TypeVar
5+ from collections .abc import Iterable , Iterator
6+ from typing import Any , Generic , TypeVar
7+
8+ if sys .version_info >= (3 , 12 ):
9+ from typing import override
10+ else :
11+ from typing_extensions import override
1012
1113from disjoint_set .utils import IdentityDict
1214
@@ -18,20 +20,20 @@ class InvalidInitialMappingError(RuntimeError):
1820
1921 def __init__ (
2022 self ,
21- msg = (
23+ msg : str = (
2224 "The mapping passed during ther DisjointSet initialization must have been wrong. "
2325 "Check that all keys are mapping to other keys and not some external values."
2426 ),
25- * args ,
26- ** kwargs ,
27+ * args : Any ,
28+ ** kwargs : Any ,
2729 ):
2830 super ().__init__ (msg , * args , ** kwargs )
2931
3032
3133class DisjointSet (Generic [T ]):
3234 """A disjoint set data structure."""
3335
34- def __init__ (self , * args , ** kwargs ) -> None :
36+ def __init__ (self , * args : Any , ** kwargs : Any ) -> None :
3537 """
3638 Disjoint set data structure.
3739
@@ -65,7 +67,8 @@ def __bool__(self) -> bool:
6567 def __getitem__ (self , element : T ) -> T :
6668 return self .find (element )
6769
68- def __eq__ (self , other : Any ) -> bool :
70+ @override
71+ def __eq__ (self , other : object ) -> bool :
6972 """
7073 Return True if both DistjoinSet structures are equivalent.
7174
@@ -76,8 +79,11 @@ def __eq__(self, other: Any) -> bool:
7679 if not isinstance (other , DisjointSet ):
7780 return False
7881
79- return {tuple (x ) for x in self .itersets ()} == {tuple (x ) for x in other .itersets ()}
82+ return {tuple (x ) for x in self .itersets ()} == {
83+ tuple (x ) for x in other .itersets ()
84+ }
8085
86+ @override
8187 def __repr__ (self ) -> str :
8288 """
8389 Print self in a reproducible way.
@@ -88,6 +94,7 @@ def __repr__(self) -> str:
8894 sets = {key : val for key , val in self }
8995 return f"{ self .__class__ .__name__ } ({ sets } )"
9096
97+ @override
9198 def __str__ (self ) -> str :
9299 return "{classname}({values})" .format (
93100 classname = self .__class__ .__name__ ,
@@ -102,7 +109,9 @@ def __iter__(self) -> Iterator[tuple[T, T]]:
102109 except RuntimeError as e :
103110 raise InvalidInitialMappingError () from e
104111
105- def itersets (self , with_canonical_elements : bool = False ) -> Iterator [set [T ] | tuple [T , set [T ]]]:
112+ def itersets (
113+ self , with_canonical_elements : bool = False
114+ ) -> Iterator [set [T ] | tuple [T , set [T ]]]:
106115 """
107116 Yield sets of connected components.
108117
@@ -114,7 +123,7 @@ def itersets(self, with_canonical_elements: bool = False) -> Iterator[set[T] | t
114123 >>> list(ds.itersets(with_canonical_elements=True))
115124 [(2, {1, 2})]
116125 """
117- element_classes : DefaultDict [T , set [T ]] = defaultdict (set )
126+ element_classes : defaultdict [T , set [T ]] = defaultdict (set )
118127 for element in self ._data :
119128 element_classes [self .find (element )].add (element )
120129
0 commit comments