networkx: annotate the linalg module

stubs/networkx/@tests/stubtest_allowlist.txt

@@ -32,6 +32,9 @@ networkx\.(utils\.)?(backends\.)?_dispatchable\.__call__
 networkx.conftest
 networkx(\..+?)?\.tests(\..+?)?
 
+# Stub-only module
+networkx\._typing
+
 # "..._DT is not present at runtime" but we don't set it in stubs, I don't understand this one
 networkx(\.algorithms)?(\.tree)?(\.mst)?\.SpanningTreeIterator\.Partition\._DT
 networkx(\.algorithms)?(\.tree)?(\.branchings)?\.ArborescenceIterator\.Partition\._DT

stubs/networkx/networkx/_typing.pyi

@@ -0,0 +1,11 @@
+from typing import TypeVar
+from typing_extensions import TypeAlias
+
+import numpy as np
+
+_G = TypeVar("_G", bound=np.generic)
+
+# numpy aliases
+Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]]
+Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_G]]
+Seed: TypeAlias = int | np.random.Generator | np.random.RandomState

stubs/networkx/networkx/linalg/algebraicconnectivity.pyi

@@ -1,28 +1,45 @@
+from typing import Literal
+
+import numpy as np
+from networkx._typing import Array1D, Seed
+from networkx.classes.graph import Graph, _Node
 from networkx.utils.backends import _dispatchable
 
 __all__ = ["algebraic_connectivity", "fiedler_vector", "spectral_ordering", "spectral_bisection"]
 
-class _PCGSolver:
-    def __init__(self, A, M) -> None: ...
-    def solve(self, B, tol): ...
-
-class _LUSolver:
-    def __init__(self, A) -> None: ...
-    def solve(self, B, tol=None): ...
-
 @_dispatchable
 def algebraic_connectivity(
-    G, weight: str = "weight", normalized: bool = False, tol: float = 1e-08, method: str = "tracemin_pcg", seed=None
-): ...
+    G: Graph[_Node],
+    weight: str | None = "weight",
+    normalized: bool = False,
+    tol: float = 1e-08,
+    method: Literal["tracemin_pcg", "tracemin_lu", "lanczos", "lobpcg"] = "tracemin_pcg",
+    seed: Seed | None = None,
+) -> float: ...
 @_dispatchable
 def fiedler_vector(
-    G, weight: str = "weight", normalized: bool = False, tol: float = 1e-08, method: str = "tracemin_pcg", seed=None
-): ...
+    G: Graph[_Node],
+    weight: str | None = "weight",
+    normalized: bool = False,
+    tol: float = 1e-08,
+    method: Literal["tracemin_pcg", "tracemin_lu", "lanczos", "lobpcg"] = "tracemin_pcg",
+    seed: Seed | None = None,
+) -> Array1D[np.float64]: ...
 @_dispatchable
 def spectral_ordering(
-    G, weight: str = "weight", normalized: bool = False, tol: float = 1e-08, method: str = "tracemin_pcg", seed=None
-): ...
+    G: Graph[_Node],
+    weight: str | None = "weight",
+    normalized: bool = False,
+    tol: float = 1e-08,
+    method: Literal["tracemin_pcg", "tracemin_lu", "lanczos", "lobpcg"] = "tracemin_pcg",
+    seed: Seed | None = None,
+) -> list[_Node]: ...
 @_dispatchable
 def spectral_bisection(
-    G, weight: str = "weight", normalized: bool = False, tol: float = 1e-08, method: str = "tracemin_pcg", seed=None
-): ...
+    G: Graph[_Node],
+    weight: str | None = "weight",
+    normalized: bool = False,
+    tol: float = 1e-08,
+    method: Literal["tracemin_pcg", "tracemin_lu", "lanczos", "lobpcg"] = "tracemin_pcg",
+    seed: Seed | None = None,
+) -> tuple[set[_Node], set[_Node]]: ...

stubs/networkx/networkx/linalg/attrmatrix.pyi

@@ -1,8 +1,41 @@
+from _typeshed import Incomplete
+from collections.abc import Collection
+from typing import Any, Literal
+
+from networkx._typing import Array2D
+from networkx.classes.graph import Graph, _Node
 from networkx.utils.backends import _dispatchable
+from numpy.typing import DTypeLike
+from scipy.sparse import lil_array  # type: ignore[import-untyped]  # pyright: ignore[reportMissingImports]
 
 __all__ = ["attr_matrix", "attr_sparse_matrix"]
 
 @_dispatchable
-def attr_matrix(G, edge_attr=None, node_attr=None, normalized: bool = False, rc_order=None, dtype=None, order=None): ...
+def attr_matrix(
+    G: Graph[_Node],
+    edge_attr: str | None = None,
+    node_attr: str | None = None,  # runtime also accepts `Callable[[_Node], object]`, but it is not documented
+    normalized: bool = False,  # runtime also accepts `Callable[[_Node, _Node], object]`, but it is not documented
+    rc_order: Collection[_Node] | None = None,
+    dtype: DTypeLike | None = None,
+    order: Literal["C", "F"] | None = None,
+    # TODO: overload on rc_order and node_attr
+    # (rc_order:[node], node_attr:None) -> 2D-array
+    # (rc_order:[any], node_attr:str) -> 2D-array
+    # (rc_order:None, node_attr:None) -> (2D-array, list[node])
+    # (rc_order:None, node_attr:str) -> (2D-array, list[any])
+) -> Array2D[Incomplete] | tuple[Array2D[Incomplete], list[_Node] | list[Any]]: ...
 @_dispatchable
-def attr_sparse_matrix(G, edge_attr=None, node_attr=None, normalized: bool = False, rc_order=None, dtype=None): ...
+def attr_sparse_matrix(
+    G: Graph[_Node],
+    edge_attr: str | None = None,
+    node_attr: str | None = None,  # runtime also accepts `Callable[[_Node], object]`, but it is not documented
+    normalized: bool = False,  # runtime also accepts `Callable[[_Node, _Node], object]`, but it is not documented
+    rc_order: Collection[_Node] | None = None,
+    dtype: DTypeLike | None = None,
+    # TODO: overload on rc_order and node_attr
+    # (rc_order:[node], node_attr:None) -> lil_array
+    # (rc_order:[any], node_attr:str) -> lil_array
+    # (rc_order:None, node_attr:None) -> (lil_array, list[node])
+    # (rc_order:None, node_attr:str) -> (lil_array, list[any])
+) -> lil_array | tuple[lil_array, list[_Node] | list[Any]]: ...

stubs/networkx/networkx/linalg/bethehessianmatrix.pyi

@@ -1,9 +1,10 @@
-from _typeshed import Incomplete
 from collections.abc import Collection
 
+from networkx.classes.graph import Graph, _Node
 from networkx.utils.backends import _dispatchable
+from scipy.sparse import csr_array  # type: ignore[import-untyped]  # pyright: ignore[reportMissingImports]
 
 __all__ = ["bethe_hessian_matrix"]
 
 @_dispatchable
-def bethe_hessian_matrix(G, r=None, nodelist: Collection[Incomplete] | None = None): ...
+def bethe_hessian_matrix(G: Graph[_Node], r: float | None = None, nodelist: Collection[_Node] | None = None) -> csr_array: ...

stubs/networkx/networkx/linalg/graphmatrix.pyi

@@ -1,11 +1,31 @@
-from _typeshed import Incomplete
-from collections.abc import Collection
+from collections.abc import Collection, Hashable
 
+from networkx.classes.graph import Graph, _Node
 from networkx.utils.backends import _dispatchable
+from numpy.typing import DTypeLike
+from scipy.sparse import csc_array, csr_array  # type: ignore[import-untyped]  # pyright: ignore[reportMissingImports]
 
 __all__ = ["incidence_matrix", "adjacency_matrix"]
 
 @_dispatchable
-def incidence_matrix(G, nodelist: Collection[Incomplete] | None = None, edgelist=None, oriented: bool = False, weight=None): ...
+def incidence_matrix(
+    G: Graph[_Node],
+    nodelist: Collection[_Node] | None = None,
+    edgelist: (
+        Collection[
+            # Requiring tuples to represent an edge might be too strict as runtime does not check the type of
+            # the collection. We can replace the tuples by `Collection[_Node | Hashable]` if people complain.
+            tuple[_Node, _Node]  # for normal graphs, this is (u, v)
+            | tuple[_Node, _Node, Hashable]  # for multigraphs, this is (u, v, key)
+        ]
+        | None
+    ) = None,
+    oriented: bool = False,
+    weight: str | None = None,
+    *,
+    dtype: DTypeLike | None = None,
+) -> csc_array: ...
 @_dispatchable
-def adjacency_matrix(G, nodelist: Collection[Incomplete] | None = None, dtype=None, weight: str = "weight"): ...
+def adjacency_matrix(
+    G: Graph[_Node], nodelist: Collection[_Node] | None = None, dtype: DTypeLike | None = None, weight: str | None = "weight"
+) -> csr_array: ...

stubs/networkx/networkx/linalg/laplacianmatrix.pyi

@@ -1,7 +1,12 @@
-from _typeshed import Incomplete
 from collections.abc import Collection
+from typing import Literal
 
+import numpy as np
+from networkx._typing import Array2D
+from networkx.classes.digraph import DiGraph
+from networkx.classes.graph import Graph, _Node
 from networkx.utils.backends import _dispatchable
+from scipy.sparse import csr_array  # type: ignore[import-untyped]  # pyright: ignore[reportMissingImports]
 
 __all__ = [
     "laplacian_matrix",
@@ -11,14 +16,24 @@ __all__ = [
 ]
 
 @_dispatchable
-def laplacian_matrix(G, nodelist: Collection[Incomplete] | None = None, weight: str = "weight"): ...
+def laplacian_matrix(G: Graph[_Node], nodelist: Collection[_Node] | None = None, weight: str | None = "weight") -> csr_array: ...
 @_dispatchable
-def normalized_laplacian_matrix(G, nodelist: Collection[Incomplete] | None = None, weight: str = "weight"): ...
+def normalized_laplacian_matrix(
+    G: Graph[_Node], nodelist: Collection[_Node] | None = None, weight: str | None = "weight"
+) -> csr_array: ...
 @_dispatchable
 def directed_laplacian_matrix(
-    G, nodelist: Collection[Incomplete] | None = None, weight: str = "weight", walk_type=None, alpha: float = 0.95
-): ...
+    G: DiGraph[_Node],
+    nodelist: Collection[_Node] | None = None,
+    weight: str | None = "weight",
+    walk_type: Literal["random", "lazy", "pagerank"] | None = None,
+    alpha: float = 0.95,
+) -> Array2D[np.float64]: ...
 @_dispatchable
 def directed_combinatorial_laplacian_matrix(
-    G, nodelist: Collection[Incomplete] | None = None, weight: str = "weight", walk_type=None, alpha: float = 0.95
-): ...
+    G: DiGraph[_Node],
+    nodelist: Collection[_Node] | None = None,
+    weight: str | None = "weight",
+    walk_type: Literal["random", "lazy", "pagerank"] | None = None,
+    alpha: float = 0.95,
+) -> Array2D[np.float64]: ...

stubs/networkx/networkx/linalg/modularitymatrix.pyi

@@ -1,11 +1,18 @@
-from _typeshed import Incomplete
 from collections.abc import Collection
 
+import numpy as np
+from networkx._typing import Array2D
+from networkx.classes.digraph import DiGraph
+from networkx.classes.graph import Graph, _Node
 from networkx.utils.backends import _dispatchable
 
 __all__ = ["modularity_matrix", "directed_modularity_matrix"]
 
 @_dispatchable
-def modularity_matrix(G, nodelist: Collection[Incomplete] | None = None, weight=None): ...
+def modularity_matrix(
+    G: Graph[_Node], nodelist: Collection[_Node] | None = None, weight: str | None = None
+) -> Array2D[np.float64]: ...
 @_dispatchable
-def directed_modularity_matrix(G, nodelist: Collection[Incomplete] | None = None, weight=None): ...
+def directed_modularity_matrix(
+    G: DiGraph[_Node], nodelist: Collection[_Node] | None = None, weight: str | None = None
+) -> Array2D[np.float64]: ...

stubs/networkx/networkx/linalg/spectrum.pyi

@@ -1,3 +1,6 @@
+import numpy as np
+from networkx._typing import Array1D
+from networkx.classes.graph import Graph, _Node
 from networkx.utils.backends import _dispatchable
 
 __all__ = [
@@ -9,12 +12,12 @@ __all__ = [
 ]
 
 @_dispatchable
-def laplacian_spectrum(G, weight: str = "weight"): ...
+def laplacian_spectrum(G: Graph[_Node], weight: str | None = "weight") -> Array1D[np.float64]: ...
 @_dispatchable
-def normalized_laplacian_spectrum(G, weight: str = "weight"): ...
+def normalized_laplacian_spectrum(G: Graph[_Node], weight: str | None = "weight") -> Array1D[np.float64]: ...
 @_dispatchable
-def adjacency_spectrum(G, weight: str = "weight"): ...
+def adjacency_spectrum(G: Graph[_Node], weight: str | None = "weight") -> Array1D[np.complex128]: ...
 @_dispatchable
-def modularity_spectrum(G): ...
+def modularity_spectrum(G: Graph[_Node]) -> Array1D[np.complex128]: ...
 @_dispatchable
-def bethe_hessian_spectrum(G, r=None): ...
+def bethe_hessian_spectrum(G: Graph[_Node], r: float | None = None) -> Array1D[np.float64]: ...