From 4bdb14ab57b6ef5937bda4fb3223e2607e05511c Mon Sep 17 00:00:00 2001 From: Ali Hamdan Date: Mon, 18 Aug 2025 21:37:56 +0200 Subject: [PATCH] networkx: complete the layout module (#14580) --- stubs/networkx/networkx/_typing.pyi | 26 +++- stubs/networkx/networkx/drawing/layout.pyi | 160 +++++++++++++-------- 2 files changed, 121 insertions(+), 65 deletions(-) diff --git a/stubs/networkx/networkx/_typing.pyi b/stubs/networkx/networkx/_typing.pyi index 6a7ab5a89..ff530c5ba 100644 --- a/stubs/networkx/networkx/_typing.pyi +++ b/stubs/networkx/networkx/_typing.pyi @@ -1,13 +1,29 @@ # Stub-only module, can't be imported at runtime. -from typing import TypeVar -from typing_extensions import TypeAlias +import sys +from collections.abc import Collection +from typing import Any, Protocol, type_check_only +from typing_extensions import TypeAlias, TypeVar import numpy as np -_G = TypeVar("_G", bound=np.generic) +_ScalarT = TypeVar("_ScalarT", bound=bool | int | float | complex | str | bytes | np.generic) +_GenericT = TypeVar("_GenericT", bound=np.generic) +_GenericT_co = TypeVar("_GenericT_co", bound=np.generic, covariant=True) +_ShapeT_co = TypeVar("_ShapeT_co", bound=tuple[int, ...], default=Any, covariant=True) # numpy aliases -Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_G]] -Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_G]] +if sys.version_info >= (3, 10): + @type_check_only + class SupportsArray(Protocol[_GenericT_co, _ShapeT_co]): + def __array__(self) -> np.ndarray[_ShapeT_co, np.dtype[_GenericT_co]]: ... + + ArrayLike1D: TypeAlias = Collection[_ScalarT] | SupportsArray[_GenericT, tuple[int]] +else: + # networkx does not support Python 3.9 but pyright still runs on 3.9 in CI + # See https://github.com/python/typeshed/issues/10722 + ArrayLike1D: TypeAlias = Collection[_ScalarT] | np.ndarray[tuple[int], np.dtype[_GenericT]] + +Array1D: TypeAlias = np.ndarray[tuple[int], np.dtype[_GenericT]] +Array2D: TypeAlias = np.ndarray[tuple[int, int], np.dtype[_GenericT]] Seed: TypeAlias = int | np.random.Generator | np.random.RandomState diff --git a/stubs/networkx/networkx/drawing/layout.pyi b/stubs/networkx/networkx/drawing/layout.pyi index 1a9a753af..1d09971ac 100644 --- a/stubs/networkx/networkx/drawing/layout.pyi +++ b/stubs/networkx/networkx/drawing/layout.pyi @@ -1,7 +1,12 @@ -from _typeshed import Incomplete +from collections.abc import Collection, Mapping +from typing import Any, Literal +from typing_extensions import TypeAlias +import numpy as np +from networkx._typing import Array1D, Array2D, ArrayLike1D, Seed +from networkx.classes.graph import Graph, _Node from networkx.utils.backends import _dispatchable -from numpy.random import RandomState +from numpy.typing import NDArray __all__ = [ "bipartite_layout", @@ -22,99 +27,134 @@ __all__ = [ "arf_layout", ] +_FloatArrayLike1D: TypeAlias = ArrayLike1D[float, np.number[Any]] # Any because we don't care about the bit base + def random_layout( - G, center=None, dim: int = 2, seed: int | RandomState | None = None, store_pos_as: str | None = None -) -> dict[Incomplete, Incomplete]: ... + G: Graph[_Node], + center: _FloatArrayLike1D | None = None, + dim: int = 2, + seed: Seed | None = None, + store_pos_as: str | None = None, +) -> dict[_Node, Array1D[np.float32]]: ... def circular_layout( - G, scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None -) -> dict[Incomplete, Incomplete]: ... + G: Graph[_Node], scale: float = 1, center: _FloatArrayLike1D | None = None, dim: int = 2, store_pos_as: str | None = None +) -> dict[_Node, Array1D[np.float64]]: ... def shell_layout( - G, nlist=None, rotate=None, scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None -) -> dict[Incomplete, Incomplete]: ... -def bipartite_layout( - G, - nodes=None, - align: str = "vertical", + G: Graph[_Node], + nlist: Collection[Collection[_Node]] | None = None, + rotate: float | None = None, scale: float = 1, - center=None, + center: _FloatArrayLike1D | None = None, + dim: int = 2, + store_pos_as: str | None = None, +) -> dict[_Node, Array1D[np.float64]]: ... +def bipartite_layout( + G: Graph[_Node], + nodes: Collection[_Node] | None = None, + align: Literal["vertical", "horizontal"] = "vertical", + scale: float = 1, + center: _FloatArrayLike1D | None = None, aspect_ratio: float = ..., store_pos_as: str | None = None, -) -> dict[Incomplete, Incomplete]: ... +) -> dict[_Node, Array1D[np.float64]]: ... def spring_layout( - G, - k=None, - pos=None, - fixed=None, + G: Graph[_Node], + k: float | None = None, + pos: Mapping[_Node, Collection[float]] | None = None, + fixed: Collection[_Node] | None = None, iterations: int = 50, threshold: float = 0.0001, - weight: str = "weight", - scale: float = 1, - center=None, + weight: str | None = "weight", + scale: float | None = 1, + center: _FloatArrayLike1D | None = None, dim: int = 2, - seed: int | RandomState | None = None, + seed: Seed | None = None, store_pos_as: str | None = None, *, - method: str = "auto", + method: Literal["auto", "force", "energy"] = "auto", gravity: float = 1.0, -) -> dict[Incomplete, Incomplete]: ... +) -> dict[_Node, Array1D[np.float64]]: ... fruchterman_reingold_layout = spring_layout def kamada_kawai_layout( - G, dist=None, pos=None, weight: str = "weight", scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None -) -> dict[Incomplete, Incomplete]: ... -def spectral_layout( - G, weight: str = "weight", scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None -) -> dict[Incomplete, Incomplete]: ... -def planar_layout( - G, scale: float = 1, center=None, dim: int = 2, store_pos_as: str | None = None -) -> dict[Incomplete, Incomplete]: ... -def spiral_layout( - G, + G: Graph[_Node], + dist: Mapping[_Node, Mapping[_Node, float]] | None = None, + pos: Mapping[_Node, Collection[float]] | None = None, + weight: str | None = "weight", scale: float = 1, - center=None, + center: _FloatArrayLike1D | None = None, + dim: int = 2, + store_pos_as: str | None = None, +) -> dict[_Node, Array1D[np.float64]]: ... +def spectral_layout( + G: Graph[_Node], + weight: str | None = "weight", + scale: float = 1, + center: _FloatArrayLike1D | None = None, + dim: int = 2, + store_pos_as: str | None = None, +) -> dict[_Node, Array1D[np.float64]]: ... +def planar_layout( + G: Graph[_Node], scale: float = 1, center: _FloatArrayLike1D | None = None, dim: int = 2, store_pos_as: str | None = None +) -> dict[_Node, Array1D[np.float64]]: ... +def spiral_layout( + G: Graph[_Node], + scale: float = 1, + center: _FloatArrayLike1D | None = None, dim: int = 2, resolution: float = 0.35, equidistant: bool = False, store_pos_as: str | None = None, -) -> dict[Incomplete, Incomplete]: ... +) -> dict[_Node, Array1D[np.float64]]: ... def multipartite_layout( - G, subset_key: str = "subset", align: str = "vertical", scale: float = 1, center=None, store_pos_as: str | None = None -) -> dict[Incomplete, Incomplete]: ... + G: Graph[_Node], + subset_key: str | Mapping[Any, Collection[_Node]] = "subset", # layers can be "any" hashable + align: Literal["vertical", "horizontal"] = "vertical", + scale: float = 1, + center: _FloatArrayLike1D | None = None, + store_pos_as: str | None = None, +) -> dict[_Node, Array1D[np.float64]]: ... def arf_layout( - G, - pos=None, + G: Graph[_Node], + pos: Mapping[_Node, Collection[float]] | None = None, scaling: float = 1, a: float = 1.1, etol: float = 1e-06, dt: float = 0.001, max_iter: int = 1000, *, - seed: int | RandomState | None = None, + seed: Seed | None = None, store_pos_as: str | None = None, -): ... +) -> dict[_Node, Array1D[np.float32]]: ... @_dispatchable def forceatlas2_layout( - G, - pos=None, + G: Graph[_Node], + pos: Mapping[_Node, Collection[float]] | None = None, *, - max_iter=100, - jitter_tolerance=1.0, - scaling_ratio=2.0, + max_iter: int = 100, + jitter_tolerance: float = 1.0, + scaling_ratio: float = 2.0, gravity: float = 1.0, - distributed_action=False, - strong_gravity=False, - node_mass=None, - node_size=None, - weight=None, - dissuade_hubs=False, - linlog=False, - seed: int | RandomState | None = None, + distributed_action: bool = False, + strong_gravity: bool = False, + node_mass: Mapping[_Node, float] | None = None, + node_size: Mapping[_Node, float] | None = None, + weight: str | None = None, + dissuade_hubs: bool = False, + linlog: bool = False, + seed: Seed | None = None, dim: int = 2, store_pos_as: str | None = None, -) -> dict[Incomplete, Incomplete]: ... -def rescale_layout(pos, scale: float = 1): ... -def rescale_layout_dict(pos, scale: float = 1): ... +) -> dict[_Node, Array1D[np.float32]]: ... +def rescale_layout(pos: NDArray[np.number[Any]], scale: float = 1) -> Array2D[np.float64]: ... # ignore the bit base +def rescale_layout_dict(pos: Mapping[_Node, Collection[float]], scale: float = 1) -> dict[_Node, Array1D[np.float64]]: ... def bfs_layout( - G, start, *, align="vertical", scale=1, center=None, store_pos_as: str | None = None -) -> dict[Incomplete, Incomplete]: ... + G: Graph[_Node], + start: _Node, + *, + align: Literal["vertical", "horizontal"] = "vertical", + scale: float = 1, + center: _FloatArrayLike1D | None = None, + store_pos_as: str | None = None, +) -> dict[_Node, Array1D[np.float64]]: ...