Skip to content

Commit 4bdb14a

Browse files
hamdanalhamdanal
authored
networkx: complete the layout module (#14580)
1 parent 540d433 commit 4bdb14a

2 files changed

Lines changed: 117 additions & 61 deletions

File tree

‎stubs/networkx/networkx/_typing.pyi‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,29 @@
11
# Stub-only module, can't be imported at runtime.
22

3-
from typing import TypeVar
4-
from typing_extensions import TypeAlias
3+
import sys
4+
from collections.abc import Collection
5+
from typing import Any, Protocol, type_check_only
6+
from typing_extensions import TypeAlias, TypeVar
57

68
import numpy as np
79

8-
_G = TypeVar("_G", bound=np.generic)
10+
_ScalarT = TypeVar("_ScalarT", bound=bool | int | float | complex | str | bytes | np.generic)
11+
_GenericT = TypeVar("_GenericT", bound=np.generic)
12+
_GenericT_co = TypeVar("_GenericT_co", bound=np.generic, covariant=True)
13+
_ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], default=Any, covariant=True)
914

1015
# numpy aliases
11-
Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]]
12-
Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_G]]
16+
if sys.version_info >= (3, 10):
17+
@type_check_only
18+
class SupportsArray(Protocol[_GenericT_co, _ShapeT_co]):
19+
def __array__(self) -> np.ndarray[_ShapeT_co, np.dtype[_GenericT_co]]: ...
20+
21+
ArrayLike1D: TypeAlias = Collection[_ScalarT] | SupportsArray[_GenericT, tuple[int]]
22+
else:
23+
# networkx does not support Python 3.9 but pyright still runs on 3.9 in CI
24+
# See https://github.com/python/typeshed/issues/10722
25+
ArrayLike1D: TypeAlias = Collection[_ScalarT] | np.ndarray[tuple[int], np.dtype[_GenericT]]
26+
27+
Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_GenericT]]
28+
Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_GenericT]]
1329
Seed: TypeAlias = int | np.random.Generator | np.random.RandomState

‎stubs/networkx/networkx/drawing/layout.pyi‎

Lines changed: 96 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,12 @@
1-
from _typeshed import Incomplete
1+
from collections.abc import Collection, Mapping
2+
from typing import Any, Literal
3+
from typing_extensions import TypeAlias
24

5+
import numpy as np
6+
from networkx._typing import Array1D, Array2D, ArrayLike1D, Seed
7+
from networkx.classes.graph import Graph, _Node
38
from networkx.utils.backends import _dispatchable
4-
from numpy.random import RandomState
9+
from numpy.typing import NDArray
510

611
__all__ = [
712
"bipartite_layout",
@@ -22,99 +27,134 @@ __all__ = [
2227
"arf_layout",
2328
]
2429

30+
_FloatArrayLike1D: TypeAlias = ArrayLike1D[float, np.number[Any]] # Any because we don't care about the bit base
31+
2532
def random_layout(
26-
G, center=None, dim: int = 2, seed: int | RandomState | None = None, store_pos_as: str | None = None
27-
) -> dict[Incomplete, Incomplete]: ...
33+
G: Graph[_Node],
34+
center: _FloatArrayLike1D | None = None,
35+
dim: int = 2,
36+
seed: Seed | None = None,
37+
store_pos_as: str | None = None,
38+
) -> dict[_Node, Array1D[np.float32]]: ...
2839
def circular_layout(
29-
G, scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None
30-
) -> dict[Incomplete, Incomplete]: ...
40+
G: Graph[_Node], scale: float = 1, center: _FloatArrayLike1D | None = None, dim: int = 2, store_pos_as: str | None = None
41+
) -> dict[_Node, Array1D[np.float64]]: ...
3142
def shell_layout(
32-
G, nlist=None, rotate=None, scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None
33-
) -> dict[Incomplete, Incomplete]: ...
43+
G: Graph[_Node],
44+
nlist: Collection[Collection[_Node]] | None = None,
45+
rotate: float | None = None,
46+
scale: float = 1,
47+
center: _FloatArrayLike1D | None = None,
48+
dim: int = 2,
49+
store_pos_as: str | None = None,
50+
) -> dict[_Node, Array1D[np.float64]]: ...
3451
def bipartite_layout(
35-
G,
36-
nodes=None,
37-
align: str = "vertical",
52+
G: Graph[_Node],
53+
nodes: Collection[_Node] | None = None,
54+
align: Literal["vertical", "horizontal"] = "vertical",
3855
scale: float = 1,
39-
center=None,
56+
center: _FloatArrayLike1D | None = None,
4057
aspect_ratio: float = ...,
4158
store_pos_as: str | None = None,
42-
) -> dict[Incomplete, Incomplete]: ...
59+
) -> dict[_Node, Array1D[np.float64]]: ...
4360
def spring_layout(
44-
G,
45-
k=None,
46-
pos=None,
47-
fixed=None,
61+
G: Graph[_Node],
62+
k: float | None = None,
63+
pos: Mapping[_Node, Collection[float]] | None = None,
64+
fixed: Collection[_Node] | None = None,
4865
iterations: int = 50,
4966
threshold: float = 0.0001,
50-
weight: str = "weight",
51-
scale: float = 1,
52-
center=None,
67+
weight: str | None = "weight",
68+
scale: float | None = 1,
69+
center: _FloatArrayLike1D | None = None,
5370
dim: int = 2,
54-
seed: int | RandomState | None = None,
71+
seed: Seed | None = None,
5572
store_pos_as: str | None = None,
5673
*,
57-
method: str = "auto",
74+
method: Literal["auto", "force", "energy"] = "auto",
5875
gravity: float = 1.0,
59-
) -> dict[Incomplete, Incomplete]: ...
76+
) -> dict[_Node, Array1D[np.float64]]: ...
6077

6178
fruchterman_reingold_layout = spring_layout
6279

6380
def kamada_kawai_layout(
64-
G, dist=None, pos=None, weight: str = "weight", scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None
65-
) -> dict[Incomplete, Incomplete]: ...
81+
G: Graph[_Node],
82+
dist: Mapping[_Node, Mapping[_Node, float]] | None = None,
83+
pos: Mapping[_Node, Collection[float]] | None = None,
84+
weight: str | None = "weight",
85+
scale: float = 1,
86+
center: _FloatArrayLike1D | None = None,
87+
dim: int = 2,
88+
store_pos_as: str | None = None,
89+
) -> dict[_Node, Array1D[np.float64]]: ...
6690
def spectral_layout(
67-
G, weight: str = "weight", scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None
68-
) -> dict[Incomplete, Incomplete]: ...
91+
G: Graph[_Node],
92+
weight: str | None = "weight",
93+
scale: float = 1,
94+
center: _FloatArrayLike1D | None = None,
95+
dim: int = 2,
96+
store_pos_as: str | None = None,
97+
) -> dict[_Node, Array1D[np.float64]]: ...
6998
def planar_layout(
70-
G, scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None
71-
) -> dict[Incomplete, Incomplete]: ...
99+
G: Graph[_Node], scale: float = 1, center: _FloatArrayLike1D | None = None, dim: int = 2, store_pos_as: str | None = None
100+
) -> dict[_Node, Array1D[np.float64]]: ...
72101
def spiral_layout(
73-
G,
102+
G: Graph[_Node],
74103
scale: float = 1,
75-
center=None,
104+
center: _FloatArrayLike1D | None = None,
76105
dim: int = 2,
77106
resolution: float = 0.35,
78107
equidistant: bool = False,
79108
store_pos_as: str | None = None,
80-
) -> dict[Incomplete, Incomplete]: ...
109+
) -> dict[_Node, Array1D[np.float64]]: ...
81110
def multipartite_layout(
82-
G, subset_key: str = "subset", align: str = "vertical", scale: float = 1, center=None, store_pos_as: str | None = None
83-
) -> dict[Incomplete, Incomplete]: ...
111+
G: Graph[_Node],
112+
subset_key: str | Mapping[Any, Collection[_Node]] = "subset", # layers can be "any" hashable
113+
align: Literal["vertical", "horizontal"] = "vertical",
114+
scale: float = 1,
115+
center: _FloatArrayLike1D | None = None,
116+
store_pos_as: str | None = None,
117+
) -> dict[_Node, Array1D[np.float64]]: ...
84118
def arf_layout(
85-
G,
86-
pos=None,
119+
G: Graph[_Node],
120+
pos: Mapping[_Node, Collection[float]] | None = None,
87121
scaling: float = 1,
88122
a: float = 1.1,
89123
etol: float = 1e-06,
90124
dt: float = 0.001,
91125
max_iter: int = 1000,
92126
*,
93-
seed: int | RandomState | None = None,
127+
seed: Seed | None = None,
94128
store_pos_as: str | None = None,
95-
): ...
129+
) -> dict[_Node, Array1D[np.float32]]: ...
96130
@_dispatchable
97131
def forceatlas2_layout(
98-
G,
99-
pos=None,
132+
G: Graph[_Node],
133+
pos: Mapping[_Node, Collection[float]] | None = None,
100134
*,
101-
max_iter=100,
102-
jitter_tolerance=1.0,
103-
scaling_ratio=2.0,
135+
max_iter: int = 100,
136+
jitter_tolerance: float = 1.0,
137+
scaling_ratio: float = 2.0,
104138
gravity: float = 1.0,
105-
distributed_action=False,
106-
strong_gravity=False,
107-
node_mass=None,
108-
node_size=None,
109-
weight=None,
110-
dissuade_hubs=False,
111-
linlog=False,
112-
seed: int | RandomState | None = None,
139+
distributed_action: bool = False,
140+
strong_gravity: bool = False,
141+
node_mass: Mapping[_Node, float] | None = None,
142+
node_size: Mapping[_Node, float] | None = None,
143+
weight: str | None = None,
144+
dissuade_hubs: bool = False,
145+
linlog: bool = False,
146+
seed: Seed | None = None,
113147
dim: int = 2,
114148
store_pos_as: str | None = None,
115-
) -> dict[Incomplete, Incomplete]: ...
116-
def rescale_layout(pos, scale: float = 1): ...
117-
def rescale_layout_dict(pos, scale: float = 1): ...
149+
) -> dict[_Node, Array1D[np.float32]]: ...
150+
def rescale_layout(pos: NDArray[np.number[Any]], scale: float = 1) -> Array2D[np.float64]: ... # ignore the bit base
151+
def rescale_layout_dict(pos: Mapping[_Node, Collection[float]], scale: float = 1) -> dict[_Node, Array1D[np.float64]]: ...
118152
def bfs_layout(
119-
G, start, *, align="vertical", scale=1, center=None, store_pos_as: str | None = None
120-
) -> dict[Incomplete, Incomplete]: ...
153+
G: Graph[_Node],
154+
start: _Node,
155+
*,
156+
align: Literal["vertical", "horizontal"] = "vertical",
157+
scale: float = 1,
158+
center: _FloatArrayLike1D | None = None,
159+
store_pos_as: str | None = None,
160+
) -> dict[_Node, Array1D[np.float64]]: ...

0 commit comments

Comments
 (0)