Fix issues with rustworkx.visit annotations (#1353)

Author: IvanIsCoding

releasenotes/notes/fix-visit-stub-eafc015adf5cada0.yaml

@@ -0,0 +1,7 @@
+---
+fixes:
+  - |
+    Fixed a bug in the discoverability of the type hints for the `rustworkx.visit` module.
+    Classes declared in the module are also now properly annotated as accepting generic types. 
+    Refer to `issue 1352 <https://github.com/Qiskit/rustworkx/issues/1352>`__ for
+    more information.

rustworkx/__init__.pyi

@@ -25,8 +25,8 @@ else:
 # rustworkx module we need to explicitly re-export every inner function from
 # rustworkx.rustworkx (the root rust module) in the form:
 # `from .rustworkx import foo as foo` so that mypy will treat `rustworkx.foo`
-# as a valid path
-import rustworkx.visit as visit
+# as a valid path.
+from . import visit as visit
 
 from .rustworkx import DAGHasCycle as DAGHasCycle
 from .rustworkx import DAGWouldCycle as DAGWouldCycle

rustworkx/visit.py

@@ -6,6 +6,10 @@
 # copyright notice, and modified files need to carry a notice indicating
 # that they have been altered from the originals.
 
+from typing import TypeVar, Generic
+
+_T = TypeVar("_T")
+
 
 class StopSearch(Exception):
     """Stop graph traversal"""
@@ -19,7 +23,7 @@ class PruneSearch(Exception):
     pass
 
 
-class BFSVisitor:
+class BFSVisitor(Generic[_T]):
     """A visitor object that is invoked at the event-points inside the
     :func:`~rustworkx.bfs_search` algorithm. By default, it performs no
     action, and should be used as a base class in order to be useful.
@@ -68,7 +72,7 @@ def black_target_edge(self, e):
         return
 
 
-class DFSVisitor:
+class DFSVisitor(Generic[_T]):
     """A visitor object that is invoked at the event-points inside the
     :func:`~rustworkx.dfs_search` algorithm. By default, it performs no
     action, and should be used as a base class in order to be useful.
@@ -119,7 +123,7 @@ def forward_or_cross_edge(self, e):
         return
 
 
-class DijkstraVisitor:
+class DijkstraVisitor(Generic[_T]):
     """A visitor object that is invoked at the event-points inside the
     :func:`~rustworkx.dijkstra_search` algorithm. By default, it performs no
     action, and should be used as a base class in order to be useful.

rustworkx/visit.pyi

@@ -9,31 +9,31 @@
 # This file contains only type annotations for PyO3 functions and classes
 # For implementation details, see visit.py
 
-from typing import Generic, TypeVar
+from typing import Any, Generic, TypeVar
 
 class StopSearch(Exception): ...
 class PruneSearch(Exception): ...
 
 _T = TypeVar("_T")
 
 class BFSVisitor(Generic[_T]):
-    def discover_vertex(self, v: int): ...
-    def finish_vertex(self, v: int): ...
-    def tree_edge(self, e: tuple[int, int, _T]): ...
-    def non_tree_edge(self, e: tuple[int, int, _T]): ...
-    def gray_target_edge(self, e: tuple[int, int, _T]): ...
-    def black_target_edge(self, e: tuple[int, int, _T]): ...
+    def discover_vertex(self, v: int) -> Any: ...
+    def finish_vertex(self, v: int) -> Any: ...
+    def tree_edge(self, e: tuple[int, int, _T]) -> Any: ...
+    def non_tree_edge(self, e: tuple[int, int, _T]) -> Any: ...
+    def gray_target_edge(self, e: tuple[int, int, _T]) -> Any: ...
+    def black_target_edge(self, e: tuple[int, int, _T]) -> Any: ...
 
 class DFSVisitor(Generic[_T]):
-    def discover_vertex(self, v: int, t: int): ...
-    def finish_vertex(self, v: int, t: int): ...
-    def tree_edge(self, e: tuple[int, int, _T]): ...
-    def back_edge(self, e: tuple[int, int, _T]): ...
-    def forward_or_cross_edge(self, e: tuple[int, int, _T]): ...
+    def discover_vertex(self, v: int, t: int) -> Any: ...
+    def finish_vertex(self, v: int, t: int) -> Any: ...
+    def tree_edge(self, e: tuple[int, int, _T]) -> Any: ...
+    def back_edge(self, e: tuple[int, int, _T]) -> Any: ...
+    def forward_or_cross_edge(self, e: tuple[int, int, _T]) -> Any: ...
 
 class DijkstraVisitor(Generic[_T]):
-    def discover_vertex(self, v: int, score: float): ...
-    def finish_vertex(self, v: int): ...
-    def examine_edge(self, edge: tuple[int, int, _T]): ...
-    def edge_relaxed(self, edge: tuple[int, int, _T]): ...
-    def edge_not_relaxed(self, edge: tuple[int, int, _T]): ...
+    def discover_vertex(self, v: int, score: float) -> Any: ...
+    def finish_vertex(self, v: int) -> Any: ...
+    def examine_edge(self, edge: tuple[int, int, _T]) -> Any: ...
+    def edge_relaxed(self, edge: tuple[int, int, _T]) -> Any: ...
+    def edge_not_relaxed(self, edge: tuple[int, int, _T]) -> Any: ...