Skip to content

Commit ea39410

Browse files
hamdanalhamdanal
authored
networkx: annotate the linalg module (#14551)
1 parent 62e28b2 commit ea39410

9 files changed

Lines changed: 151 additions & 39 deletions

File tree

‎stubs/networkx/@tests/stubtest_allowlist.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,9 @@ networkx\.(utils\.)?(backends\.)?_dispatchable\.__call__
3131
networkx.conftest
3232
networkx(\..+?)?\.tests(\..+?)?
3333

34+
# Stub-only module
35+
networkx\._typing
36+
3437
# "..._DT is not present at runtime" but we don't set it in stubs, I don't understand this one
3538
networkx(\.algorithms)?(\.tree)?(\.mst)?\.SpanningTreeIterator\.Partition\._DT
3639
networkx(\.algorithms)?(\.tree)?(\.branchings)?\.ArborescenceIterator\.Partition\._DT

‎stubs/networkx/networkx/_typing.pyi‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# Stub-only module, can't be imported at runtime.
2+
3+
from typing import TypeVar
4+
from typing_extensions import TypeAlias
5+
6+
import numpy as np
7+
8+
_G = TypeVar("_G", bound=np.generic)
9+
10+
# numpy aliases
11+
Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]]
12+
Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_G]]
13+
Seed: TypeAlias = int | np.random.Generator | np.random.RandomState

‎stubs/networkx/networkx/linalg/algebraicconnectivity.pyi‎

Lines changed: 33 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,28 +1,45 @@
1+
from typing import Literal
2+
3+
import numpy as np
4+
from networkx._typing import Array1D, Seed
5+
from networkx.classes.graph import Graph, _Node
16
from networkx.utils.backends import _dispatchable
27

38
__all__ = ["algebraic_connectivity", "fiedler_vector", "spectral_ordering", "spectral_bisection"]
49

5-
class _PCGSolver:
6-
def __init__(self, A, M) -> None: ...
7-
def solve(self, B, tol): ...
8-
9-
class _LUSolver:
10-
def __init__(self, A) -> None: ...
11-
def solve(self, B, tol=None): ...
12-
1310
@_dispatchable
1411
def algebraic_connectivity(
15-
G, weight: str = "weight", normalized: bool = False, tol: float = 1e-08, method: str = "tracemin_pcg", seed=None
16-
): ...
12+
G: Graph[_Node],
13+
weight: str | None = "weight",
14+
normalized: bool = False,
15+
tol: float = 1e-08,
16+
method: Literal["tracemin_pcg", "tracemin_lu", "lanczos", "lobpcg"] = "tracemin_pcg",
17+
seed: Seed | None = None,
18+
) -> float: ...
1719
@_dispatchable
1820
def fiedler_vector(
19-
G, weight: str = "weight", normalized: bool = False, tol: float = 1e-08, method: str = "tracemin_pcg", seed=None
20-
): ...
21+
G: Graph[_Node],
22+
weight: str | None = "weight",
23+
normalized: bool = False,
24+
tol: float = 1e-08,
25+
method: Literal["tracemin_pcg", "tracemin_lu", "lanczos", "lobpcg"] = "tracemin_pcg",
26+
seed: Seed | None = None,
27+
) -> Array1D[np.float64]: ...
2128
@_dispatchable
2229
def spectral_ordering(
23-
G, weight: str = "weight", normalized: bool = False, tol: float = 1e-08, method: str = "tracemin_pcg", seed=None
24-
): ...
30+
G: Graph[_Node],
31+
weight: str | None = "weight",
32+
normalized: bool = False,
33+
tol: float = 1e-08,
34+
method: Literal["tracemin_pcg", "tracemin_lu", "lanczos", "lobpcg"] = "tracemin_pcg",
35+
seed: Seed | None = None,
36+
) -> list[_Node]: ...
2537
@_dispatchable
2638
def spectral_bisection(
27-
G, weight: str = "weight", normalized: bool = False, tol: float = 1e-08, method: str = "tracemin_pcg", seed=None
28-
): ...
39+
G: Graph[_Node],
40+
weight: str | None = "weight",
41+
normalized: bool = False,
42+
tol: float = 1e-08,
43+
method: Literal["tracemin_pcg", "tracemin_lu", "lanczos", "lobpcg"] = "tracemin_pcg",
44+
seed: Seed | None = None,
45+
) -> tuple[set[_Node], set[_Node]]: ...

‎stubs/networkx/networkx/linalg/attrmatrix.pyi‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,41 @@
1+
from _typeshed import Incomplete
2+
from collections.abc import Collection
3+
from typing import Any, Literal
4+
5+
from networkx._typing import Array2D
6+
from networkx.classes.graph import Graph, _Node
17
from networkx.utils.backends import _dispatchable
8+
from numpy.typing import DTypeLike
9+
from scipy.sparse import lil_array # type: ignore[import-untyped] # pyright: ignore[reportMissingImports]
210

311
__all__ = ["attr_matrix", "attr_sparse_matrix"]
412

513
@_dispatchable
6-
def attr_matrix(G, edge_attr=None, node_attr=None, normalized: bool = False, rc_order=None, dtype=None, order=None): ...
14+
def attr_matrix(
15+
G: Graph[_Node],
16+
edge_attr: str | None = None,
17+
node_attr: str | None = None, # runtime also accepts `Callable[[_Node], object]`, but it is not documented
18+
normalized: bool = False, # runtime also accepts `Callable[[_Node, _Node], object]`, but it is not documented
19+
rc_order: Collection[_Node] | None = None,
20+
dtype: DTypeLike | None = None,
21+
order: Literal["C", "F"] | None = None,
22+
# TODO: overload on rc_order and node_attr
23+
# (rc_order:[node], node_attr:None) -> 2D-array
24+
# (rc_order:[any], node_attr:str) -> 2D-array
25+
# (rc_order:None, node_attr:None) -> (2D-array, list[node])
26+
# (rc_order:None, node_attr:str) -> (2D-array, list[any])
27+
) -> Array2D[Incomplete] | tuple[Array2D[Incomplete], list[_Node] | list[Any]]: ...
728
@_dispatchable
8-
def attr_sparse_matrix(G, edge_attr=None, node_attr=None, normalized: bool = False, rc_order=None, dtype=None): ...
29+
def attr_sparse_matrix(
30+
G: Graph[_Node],
31+
edge_attr: str | None = None,
32+
node_attr: str | None = None, # runtime also accepts `Callable[[_Node], object]`, but it is not documented
33+
normalized: bool = False, # runtime also accepts `Callable[[_Node, _Node], object]`, but it is not documented
34+
rc_order: Collection[_Node] | None = None,
35+
dtype: DTypeLike | None = None,
36+
# TODO: overload on rc_order and node_attr
37+
# (rc_order:[node], node_attr:None) -> lil_array
38+
# (rc_order:[any], node_attr:str) -> lil_array
39+
# (rc_order:None, node_attr:None) -> (lil_array, list[node])
40+
# (rc_order:None, node_attr:str) -> (lil_array, list[any])
41+
) -> lil_array | tuple[lil_array, list[_Node] | list[Any]]: ...

‎stubs/networkx/networkx/linalg/bethehessianmatrix.pyi‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
1-
from _typeshed import Incomplete
21
from collections.abc import Collection
32

3+
from networkx.classes.graph import Graph, _Node
44
from networkx.utils.backends import _dispatchable
5+
from scipy.sparse import csr_array # type: ignore[import-untyped] # pyright: ignore[reportMissingImports]
56

67
__all__ = ["bethe_hessian_matrix"]
78

89
@_dispatchable
9-
def bethe_hessian_matrix(G, r=None, nodelist: Collection[Incomplete] | None = None): ...
10+
def bethe_hessian_matrix(G: Graph[_Node], r: float | None = None, nodelist: Collection[_Node] | None = None) -> csr_array: ...

‎stubs/networkx/networkx/linalg/graphmatrix.pyi‎

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,31 @@
1-
from _typeshed import Incomplete
2-
from collections.abc import Collection
1+
from collections.abc import Collection, Hashable
32

3+
from networkx.classes.graph import Graph, _Node
44
from networkx.utils.backends import _dispatchable
5+
from numpy.typing import DTypeLike
6+
from scipy.sparse import csc_array, csr_array # type: ignore[import-untyped] # pyright: ignore[reportMissingImports]
57

68
__all__ = ["incidence_matrix", "adjacency_matrix"]
79

810
@_dispatchable
9-
def incidence_matrix(G, nodelist: Collection[Incomplete] | None = None, edgelist=None, oriented: bool = False, weight=None): ...
11+
def incidence_matrix(
12+
G: Graph[_Node],
13+
nodelist: Collection[_Node] | None = None,
14+
edgelist: (
15+
Collection[
16+
# Requiring tuples to represent an edge might be too strict as runtime does not check the type of
17+
# the collection. We can replace the tuples by `Collection[_Node | Hashable]` if people complain.
18+
tuple[_Node, _Node] # for normal graphs, this is (u, v)
19+
| tuple[_Node, _Node, Hashable] # for multigraphs, this is (u, v, key)
20+
]
21+
| None
22+
) = None,
23+
oriented: bool = False,
24+
weight: str | None = None,
25+
*,
26+
dtype: DTypeLike | None = None,
27+
) -> csc_array: ...
1028
@_dispatchable
11-
def adjacency_matrix(G, nodelist: Collection[Incomplete] | None = None, dtype=None, weight: str = "weight"): ...
29+
def adjacency_matrix(
30+
G: Graph[_Node], nodelist: Collection[_Node] | None = None, dtype: DTypeLike | None = None, weight: str | None = "weight"
31+
) -> csr_array: ...

‎stubs/networkx/networkx/linalg/laplacianmatrix.pyi‎

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
1-
from _typeshed import Incomplete
21
from collections.abc import Collection
2+
from typing import Literal
33

4+
import numpy as np
5+
from networkx._typing import Array2D
6+
from networkx.classes.digraph import DiGraph
7+
from networkx.classes.graph import Graph, _Node
48
from networkx.utils.backends import _dispatchable
9+
from scipy.sparse import csr_array # type: ignore[import-untyped] # pyright: ignore[reportMissingImports]
510

611
__all__ = [
712
"laplacian_matrix",
@@ -11,14 +16,24 @@ __all__ = [
1116
]
1217

1318
@_dispatchable
14-
def laplacian_matrix(G, nodelist: Collection[Incomplete] | None = None, weight: str = "weight"): ...
19+
def laplacian_matrix(G: Graph[_Node], nodelist: Collection[_Node] | None = None, weight: str | None = "weight") -> csr_array: ...
1520
@_dispatchable
16-
def normalized_laplacian_matrix(G, nodelist: Collection[Incomplete] | None = None, weight: str = "weight"): ...
21+
def normalized_laplacian_matrix(
22+
G: Graph[_Node], nodelist: Collection[_Node] | None = None, weight: str | None = "weight"
23+
) -> csr_array: ...
1724
@_dispatchable
1825
def directed_laplacian_matrix(
19-
G, nodelist: Collection[Incomplete] | None = None, weight: str = "weight", walk_type=None, alpha: float = 0.95
20-
): ...
26+
G: DiGraph[_Node],
27+
nodelist: Collection[_Node] | None = None,
28+
weight: str | None = "weight",
29+
walk_type: Literal["random", "lazy", "pagerank"] | None = None,
30+
alpha: float = 0.95,
31+
) -> Array2D[np.float64]: ...
2132
@_dispatchable
2233
def directed_combinatorial_laplacian_matrix(
23-
G, nodelist: Collection[Incomplete] | None = None, weight: str = "weight", walk_type=None, alpha: float = 0.95
24-
): ...
34+
G: DiGraph[_Node],
35+
nodelist: Collection[_Node] | None = None,
36+
weight: str | None = "weight",
37+
walk_type: Literal["random", "lazy", "pagerank"] | None = None,
38+
alpha: float = 0.95,
39+
) -> Array2D[np.float64]: ...

‎stubs/networkx/networkx/linalg/modularitymatrix.pyi‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,18 @@
1-
from _typeshed import Incomplete
21
from collections.abc import Collection
32

3+
import numpy as np
4+
from networkx._typing import Array2D
5+
from networkx.classes.digraph import DiGraph
6+
from networkx.classes.graph import Graph, _Node
47
from networkx.utils.backends import _dispatchable
58

69
__all__ = ["modularity_matrix", "directed_modularity_matrix"]
710

811
@_dispatchable
9-
def modularity_matrix(G, nodelist: Collection[Incomplete] | None = None, weight=None): ...
12+
def modularity_matrix(
13+
G: Graph[_Node], nodelist: Collection[_Node] | None = None, weight: str | None = None
14+
) -> Array2D[np.float64]: ...
1015
@_dispatchable
11-
def directed_modularity_matrix(G, nodelist: Collection[Incomplete] | None = None, weight=None): ...
16+
def directed_modularity_matrix(
17+
G: DiGraph[_Node], nodelist: Collection[_Node] | None = None, weight: str | None = None
18+
) -> Array2D[np.float64]: ...

‎stubs/networkx/networkx/linalg/spectrum.pyi‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,6 @@
1+
import numpy as np
2+
from networkx._typing import Array1D
3+
from networkx.classes.graph import Graph, _Node
14
from networkx.utils.backends import _dispatchable
25

36
__all__ = [
@@ -9,12 +12,12 @@ __all__ = [
912
]
1013

1114
@_dispatchable
12-
def laplacian_spectrum(G, weight: str = "weight"): ...
15+
def laplacian_spectrum(G: Graph[_Node], weight: str | None = "weight") -> Array1D[np.float64]: ...
1316
@_dispatchable
14-
def normalized_laplacian_spectrum(G, weight: str = "weight"): ...
17+
def normalized_laplacian_spectrum(G: Graph[_Node], weight: str | None = "weight") -> Array1D[np.float64]: ...
1518
@_dispatchable
16-
def adjacency_spectrum(G, weight: str = "weight"): ...
19+
def adjacency_spectrum(G: Graph[_Node], weight: str | None = "weight") -> Array1D[np.complex128]: ...
1720
@_dispatchable
18-
def modularity_spectrum(G): ...
21+
def modularity_spectrum(G: Graph[_Node]) -> Array1D[np.complex128]: ...
1922
@_dispatchable
20-
def bethe_hessian_spectrum(G, r=None): ...
23+
def bethe_hessian_spectrum(G: Graph[_Node], r: float | None = None) -> Array1D[np.float64]: ...

0 commit comments

Comments
 (0)