[networkx] add type hints for networkx.utils.union_find.UnionFind

Author: cj81499

stubs/networkx/networkx/utils/union_find.pyi

@@ -1,11 +1,13 @@
-from _typeshed import Incomplete
-from collections.abc import Generator, Iterator
+from collections.abc import Generator, Iterable, Iterator, Mapping
+from typing import Generic, TypeVar
 
-class UnionFind:
-    parents: Incomplete
-    weights: Incomplete
-    def __init__(self, elements=None) -> None: ...
-    def __getitem__(self, object): ...
-    def __iter__(self) -> Iterator[Incomplete]: ...
-    def to_sets(self) -> Generator[Incomplete, Incomplete, None]: ...
-    def union(self, *objects): ...
+_T = TypeVar("_T")
+
+class UnionFind(Generic[_T]):
+    parents: Mapping[_T, _T]
+    weights: Mapping[_T, int]
+    def __init__(self, elements: Iterable[_T] | None = None) -> None: ...
+    def __getitem__(self, object: _T) -> _T: ...
+    def __iter__(self) -> Iterator[_T]: ...
+    def to_sets(self) -> Generator[set[_T]]: ...
+    def union(self, *objects: _T) -> None: ...