[networkx] Accept any node/edge data types in graph parameters outside algorithms

Graph parameters annotated as `Graph[_Node]` only accept the default
`dict[str, Any]` node and edge data, so graphs with other `Mapping` data
types are rejected by functions that never look at the data. Type them as
`Graph[_Node, _NodeData, _EdgeData]` in `classes.function`, `convert`,
`convert_matrix`, `generators`, `linalg`, `readwrite` and `utils`, as
already done for `classes` and `drawing`. Return types of graph-building
functions are unchanged; `edge_subgraph()` and `restricted_view()` now
return the input graph's type.

Part of #16365.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

Author: ekanshul

stubs/networkx/@tests/test_cases/check_graph_data_types.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping
+from typing import Any
+
+from typing_extensions import assert_type
+
+import networkx as nx
+from networkx.utils.rcm import reverse_cuthill_mckee_ordering
+
+
+class NodeData(Mapping[str, Any]):
+    def __getitem__(self, key: str) -> Any: ...
+    def __iter__(self) -> Iterator[str]: ...
+    def __len__(self) -> int: ...
+
+
+# Functions that only read a graph accept any node/edge data types that satisfy
+# the `Mapping[str, Any]` bound, not just the `dict[str, Any]` default.
+G = nx.Graph[int, NodeData, dict[str, Any]]()
+assert_type(nx.degree_histogram(G), list[int])
+assert_type(nx.to_dict_of_lists(G), dict[int, list[int]])
+assert_type(nx.is_weighted(G), bool)
+nx.number_of_nodes(G)
+nx.write_gml(G, "graph.gml")
+nx.generate_adjlist(G)
+nx.node_link_data(G)
+nx.adjacency_matrix(G)
+nx.laplacian_spectrum(G)
+reverse_cuthill_mckee_ordering(G)
+
+D = nx.DiGraph[str, NodeData, NodeData]()
+nx.number_of_edges(D)
+nx.write_edgelist(D, "graph.edgelist")
+nx.directed_laplacian_matrix(D)
+
+# Views keep the data types of the graph they come from.
+assert_type(nx.edge_subgraph(G, [(1, 2)]), nx.Graph[int, NodeData, dict[str, Any]])
+assert_type(nx.restricted_view(G, [1], [(1, 2)]), nx.Graph[int, NodeData, dict[str, Any]])
+
+# The default data types work as before.
+H = nx.Graph[str]()
+assert_type(nx.to_dict_of_lists(H), dict[str, list[str]])
+nx.write_gml(H, "graph.gml")

stubs/networkx/networkx/classes/function.pyi

@@ -1,6 +1,6 @@
 from _typeshed import Incomplete, SupportsItems, SupportsKeysAndGetItem, Unused
 from collections.abc import Callable, Collection, Generator, Hashable, Iterable, Iterator
-from typing import Literal, TypeVar, overload
+from typing import Any, Literal, TypeVar, overload
 
 from networkx import _dispatchable
 from networkx.algorithms.planarity import PlanarEmbedding
@@ -53,34 +53,34 @@ __all__ = [
 
 _U = TypeVar("_U")
 
-def nodes(G: Graph[_Node]): ...
-def edges(G: Graph[_Node], nbunch=None): ...
-def degree(G: Graph[_Node], nbunch=None, weight=None): ...
-def neighbors(G: Graph[_Node], n): ...
-def number_of_nodes(G: Graph[_Node]): ...
-def number_of_edges(G: Graph[_Node]): ...
-def density(G: Graph[_Node]): ...
-def degree_histogram(G: Graph[_Node]) -> list[int]: ...
-
+def nodes(G: Graph[_Node, _NodeData, _EdgeData]): ...
+def edges(G: Graph[_Node, _NodeData, _EdgeData], nbunch=None): ...
+def degree(G: Graph[_Node, _NodeData, _EdgeData], nbunch=None, weight=None): ...
+def neighbors(G: Graph[_Node, _NodeData, _EdgeData], n): ...
+def number_of_nodes(G: Graph[_Node, _NodeData, _EdgeData]): ...
+def number_of_edges(G: Graph[_Node, _NodeData, _EdgeData]): ...
+def density(G: Graph[_Node, _NodeData, _EdgeData]): ...
+def degree_histogram(G: Graph[_Node, _NodeData, _EdgeData]) -> list[int]: ...
 @overload
 def is_directed(G: PlanarEmbedding[Hashable]) -> Literal[False]: ...  # type: ignore[misc] # Incompatible return types
 @overload
 def is_directed(G: DiGraph[Hashable]) -> Literal[True]: ...  # type: ignore[misc] # Incompatible return types
 @overload
 def is_directed(G: Graph[Hashable]) -> Literal[False]: ...
-
-def freeze(G: Graph[_Node]): ...
+def freeze(G: Graph[_Node, _NodeData, _EdgeData]): ...
 def is_frozen(G: Graph[Incomplete]) -> bool: ...
 def add_star(G_to_add_to: Graph[Incomplete], nodes_for_star: Iterable[Incomplete], **attr) -> None: ...
 def add_path(G_to_add_to: Graph[Incomplete], nodes_for_path: Iterable[Incomplete], **attr) -> None: ...
 def add_cycle(G_to_add_to: Graph[Incomplete], nodes_for_cycle: Iterable[Incomplete], **attr) -> None: ...
-def subgraph(G: Graph[_Node], nbunch: Iterable[Incomplete]): ...
+def subgraph(G: Graph[_Node, _NodeData, _EdgeData], nbunch: Iterable[Incomplete]): ...
 def induced_subgraph(G: Graph[_Node, _NodeData, _EdgeData], nbunch: _NBunch[_Node]) -> Graph[_Node, _NodeData, _EdgeData]: ...
-def edge_subgraph(G: Graph[_Node], edges: Iterable[Incomplete]) -> Graph[Incomplete]: ...
-def restricted_view(G: Graph[_Node], nodes: Iterable[Incomplete], edges: Iterable[Incomplete]) -> Graph[Incomplete]: ...
+def edge_subgraph(G: Graph[_Node, _NodeData, _EdgeData], edges: Iterable[Incomplete]) -> Graph[_Node, _NodeData, _EdgeData]: ...
+def restricted_view(
+    G: Graph[_Node, _NodeData, _EdgeData], nodes: Iterable[Incomplete], edges: Iterable[Incomplete]
+) -> Graph[_Node, _NodeData, _EdgeData]: ...
 def to_directed(graph): ...
 def to_undirected(graph): ...
-def create_empty_copy(G: Graph[_Node], with_data: bool = True): ...
+def create_empty_copy(G: Graph[_Node, _NodeData, _EdgeData], with_data: bool = True): ...
 
 # incomplete: Can "Any scalar value" be enforced?
 @overload
@@ -94,22 +94,20 @@ def set_node_attributes(
 ) -> None: ...
 @overload
 def set_node_attributes(
-    G: Graph[_Node],
+    G: Graph[_Node, _NodeData, _EdgeData],
     values: SupportsItems[_Node, SupportsKeysAndGetItem[Incomplete, Incomplete] | Iterable[tuple[Incomplete, Incomplete]]],
     name: None = None,
     *,
     backend=None,
     **backend_kwargs,
 ) -> None: ...
-
 @_dispatchable
-def get_node_attributes(G: Graph[_Node], name: str, default=None) -> dict[_Node, Incomplete]: ...
+def get_node_attributes(G: Graph[_Node, _NodeData, _EdgeData], name: str, default=None) -> dict[_Node, Incomplete]: ...
 @_dispatchable
-def remove_node_attributes(G: Graph[_Node], *attr_names, nbunch=None) -> None: ...
-
+def remove_node_attributes(G: Graph[_Node, _NodeData, _EdgeData], *attr_names, nbunch=None) -> None: ...
 @overload
 def set_edge_attributes(
-    G: Graph[_Node],
+    G: Graph[_Node, _NodeData, _EdgeData],
     values: SupportsItems[tuple[_Node, _Node], Incomplete],
     name: str,
     *,
@@ -118,7 +116,7 @@ def set_edge_attributes(
 ) -> None: ...
 @overload
 def set_edge_attributes(
-    G: MultiGraph[_Node],
+    G: MultiGraph[_Node, _NodeData, _EdgeData],
     values: dict[tuple[_Node, _Node, Incomplete], Incomplete],
     name: str,
     *,
@@ -129,54 +127,60 @@ def set_edge_attributes(
 def set_edge_attributes(
     G: Graph[Hashable], values, name: None = None, *, backend: str | None = None, **backend_kwargs
 ) -> None: ...
-
 @_dispatchable
-def get_edge_attributes(G: Graph[_Node], name: str, default=None) -> dict[tuple[_Node, _Node], Incomplete]: ...
+def get_edge_attributes(
+    G: Graph[_Node, _NodeData, _EdgeData], name: str, default=None
+) -> dict[tuple[_Node, _Node], Incomplete]: ...
 @_dispatchable
-def remove_edge_attributes(G: Graph[_Node], *attr_names, ebunch=None) -> None: ...
-def all_neighbors(graph: Graph[_Node], node: _Node) -> Iterator[_Node]: ...
-def non_neighbors(graph: Graph[_Node], node: _Node) -> Generator[_Node]: ...
-def non_edges(graph: Graph[_Node]) -> Generator[tuple[_Node, _Node]]: ...
-def common_neighbors(G: Graph[_Node], u: _Node, v: _Node) -> Generator[_Node]: ...
+def remove_edge_attributes(G: Graph[_Node, _NodeData, _EdgeData], *attr_names, ebunch=None) -> None: ...
+def all_neighbors(graph: Graph[_Node, _NodeData, _EdgeData], node: _Node) -> Iterator[_Node]: ...
+def non_neighbors(graph: Graph[_Node, _NodeData, _EdgeData], node: _Node) -> Generator[_Node]: ...
+def non_edges(graph: Graph[_Node, _NodeData, _EdgeData]) -> Generator[tuple[_Node, _Node]]: ...
+def common_neighbors(G: Graph[_Node, _NodeData, _EdgeData], u: _Node, v: _Node) -> Generator[_Node]: ...
 @_dispatchable
-def is_weighted(G: Graph[_Node], edge: tuple[_Node, _Node] | None = None, weight: str = "weight") -> bool: ...
+def is_weighted(
+    G: Graph[_Node, _NodeData, _EdgeData], edge: tuple[_Node, _Node] | None = None, weight: str = "weight"
+) -> bool: ...
 @_dispatchable
-def is_negatively_weighted(G: Graph[_Node], edge: tuple[_Node, _Node] | None = None, weight: str = "weight") -> bool: ...
+def is_negatively_weighted(
+    G: Graph[_Node, _NodeData, _EdgeData], edge: tuple[_Node, _Node] | None = None, weight: str = "weight"
+) -> bool: ...
 @_dispatchable
 def is_empty(G: Graph[Hashable]) -> bool: ...
-def nodes_with_selfloops(G: Graph[_Node]) -> Generator[_Node]: ...
-
+def nodes_with_selfloops(G: Graph[_Node, _NodeData, _EdgeData]) -> Generator[_Node]: ...
 @overload
 def selfloop_edges(
-    G: Graph[_Node], data: Literal[False] = False, keys: Literal[False] = False, default=None
+    G: Graph[_Node, _NodeData, _EdgeData], data: Literal[False] = False, keys: Literal[False] = False, default=None
 ) -> Generator[tuple[_Node, _Node]]: ...
 @overload
 def selfloop_edges(
     G: Graph[_Node, _NodeData, _EdgeData], data: Literal[True], keys: Literal[False] = False, default=None
 ) -> Generator[tuple[_Node, _Node, _EdgeData]]: ...
 @overload
 def selfloop_edges(
-    G: Graph[_Node], data: str, keys: Literal[False] = False, default: _U | None = None
+    G: Graph[_Node, Any, Any], data: str, keys: Literal[False] = False, default: _U | None = None
 ) -> Generator[tuple[_Node, _Node, _U]]: ...
 @overload
 def selfloop_edges(
-    G: Graph[_Node], data: Literal[False], keys: Literal[True], default=None
+    G: Graph[_Node, _NodeData, _EdgeData], data: Literal[False], keys: Literal[True], default=None
 ) -> Generator[tuple[_Node, _Node, int]]: ...
 @overload
 def selfloop_edges(
-    G: Graph[_Node], data: Literal[False] = False, *, keys: Literal[True], default=None
+    G: Graph[_Node, _NodeData, _EdgeData], data: Literal[False] = False, *, keys: Literal[True], default=None
 ) -> Generator[tuple[_Node, _Node, int]]: ...
 @overload
 def selfloop_edges(
     G: Graph[_Node, _NodeData, _EdgeData], data: Literal[True], keys: Literal[True], default=None
 ) -> Generator[tuple[_Node, _Node, int, _EdgeData]]: ...
 @overload
 def selfloop_edges(
-    G: Graph[_Node], data: str, keys: Literal[True], default: _U | None = None
+    G: Graph[_Node, Any, Any], data: str, keys: Literal[True], default: _U | None = None
 ) -> Generator[tuple[_Node, _Node, int, _U]]: ...
-
 @_dispatchable
 def number_of_selfloops(G: Graph[Hashable]) -> int: ...
-def is_path(G: Graph[_Node], path: Iterable[Incomplete]) -> bool: ...
-def path_weight(G: Graph[_Node], path: Collection[Incomplete], weight: str) -> int: ...
-def describe(G: Graph[_Node], describe_hook: Callable[[Graph[_Node]], dict[str, Incomplete]] | None = None) -> None: ...
+def is_path(G: Graph[_Node, _NodeData, _EdgeData], path: Iterable[Incomplete]) -> bool: ...
+def path_weight(G: Graph[_Node, _NodeData, _EdgeData], path: Collection[Incomplete], weight: str) -> int: ...
+def describe(
+    G: Graph[_Node, _NodeData, _EdgeData],
+    describe_hook: Callable[[Graph[_Node, _NodeData, _EdgeData]], dict[str, Incomplete]] | None = None,
+) -> None: ...

stubs/networkx/networkx/convert.pyi

@@ -20,13 +20,15 @@ def to_networkx_graph(
     multigraph_input: bool = False,
 ) -> Graph[_Node, _NodeData, _EdgeData]: ...
 @_dispatchable
-def to_dict_of_lists(G: Graph[_Node], nodelist: Collection[_Node] | None = None) -> dict[_Node, list[_Node]]: ...
+def to_dict_of_lists(
+    G: Graph[_Node, _NodeData, _EdgeData], nodelist: Collection[_Node] | None = None
+) -> dict[_Node, list[_Node]]: ...
 @_dispatchable
 def from_dict_of_lists(
     d: dict[_Node, Iterable[_Node]], create_using: Graph[Incomplete] | type[Graph[Incomplete]] | None = None
 ) -> Graph[_Node]: ...
 def to_dict_of_dicts(
-    G: Graph[_Node], nodelist: Collection[_Node] | None = None, edge_data: float | None = None
+    G: Graph[_Node, _NodeData, _EdgeData], nodelist: Collection[_Node] | None = None, edge_data: float | None = None
 ) -> dict[Incomplete, Incomplete]: ...
 @_dispatchable
 def from_dict_of_dicts(
@@ -35,7 +37,7 @@ def from_dict_of_dicts(
     multigraph_input: bool = False,
 ) -> Graph[Incomplete]: ...
 @_dispatchable
-def to_edgelist(G: Graph[_Node], nodelist: Collection[_Node] | None = None): ...
+def to_edgelist(G: Graph[_Node, _NodeData, _EdgeData], nodelist: Collection[_Node] | None = None): ...
 @_dispatchable
 def from_edgelist(
     edgelist: Iterable[Incomplete], create_using: Graph[Incomplete] | type[Graph[Incomplete]] | None = None

stubs/networkx/networkx/convert_matrix.pyi

@@ -4,7 +4,7 @@ from typing import Literal, TypeAlias, TypeVar, overload
 
 import numpy
 import numpy as np
-from networkx.classes.graph import Graph, _Node
+from networkx.classes.graph import Graph, _EdgeData, _Node, _NodeData
 from networkx.utils.backends import _dispatchable
 
 # stub_uploader won't allow pandas-stubs in the requires field https://github.com/typeshed-internal/stub_uploader/issues/90
@@ -30,30 +30,27 @@ __all__ = [
 
 @_dispatchable
 def to_pandas_adjacency(
-    G: Graph[_Node],
+    G: Graph[_Node, _NodeData, _EdgeData],
     nodelist: _Axes[_Node] | None = None,
     dtype: numpy.dtype[Incomplete] | None = None,
     order: numpy._OrderCF = None,
     multigraph_weight: Callable[[list[float]], float] = ...,
     weight: str = "weight",
     nonedge: float = 0.0,
 ) -> _DataFrame: ...
-
 @overload
 def from_pandas_adjacency(df: _DataFrame, create_using: type[_G]) -> _G: ...
 @overload
 def from_pandas_adjacency(df: _DataFrame, create_using: None = None) -> Graph[Incomplete]: ...
-
 @_dispatchable
 def to_pandas_edgelist(
-    G: Graph[_Node],
+    G: Graph[_Node, _NodeData, _EdgeData],
     source: str | int = "source",
     target: str | int = "target",
     nodelist: Iterable[_Node] | None = None,
     dtype: _ExtensionDtype | None = None,
     edge_key: str | int | None = None,
 ) -> _DataFrame: ...
-
 @overload
 def from_pandas_edgelist(
     df: _DataFrame,
@@ -82,10 +79,9 @@ def from_pandas_edgelist(
     create_using: None = None,
     edge_key: str | None = None,
 ) -> Graph[Incomplete]: ...
-
 @_dispatchable
 def to_scipy_sparse_array(
-    G: Graph[_Node],
+    G: Graph[_Node, _NodeData, _EdgeData],
     nodelist: Collection[_Node] | None = None,
     dtype: np.dtype[Incomplete] | None = None,
     weight: str | None = "weight",
@@ -100,15 +96,14 @@ def from_scipy_sparse_array(
 ): ...
 @_dispatchable
 def to_numpy_array(
-    G: Graph[_Node],
+    G: Graph[_Node, _NodeData, _EdgeData],
     nodelist: Collection[_Node] | None = None,
     dtype: numpy.dtype[Incomplete] | None = None,
     order: numpy._OrderCF = None,
     multigraph_weight: Callable[[list[float]], float] = ...,
     weight: str = "weight",
     nonedge: float = 0.0,
 ) -> numpy.ndarray[Incomplete, numpy.dtype[Incomplete]]: ...
-
 @overload
 def from_numpy_array(
     A: numpy.ndarray[Incomplete, Incomplete], parallel_edges: bool = False, create_using: None = None

stubs/networkx/networkx/generators/ego.pyi

@@ -1,7 +1,9 @@
-from networkx.classes.graph import Graph, _Node
+from networkx.classes.graph import Graph, _EdgeData, _Node, _NodeData
 from networkx.utils.backends import _dispatchable
 
 __all__ = ["ego_graph"]
 
 @_dispatchable
-def ego_graph(G: Graph[_Node], n, radius: float = 1, center: bool = True, undirected: bool = False, distance=None): ...
+def ego_graph(
+    G: Graph[_Node, _NodeData, _EdgeData], n, radius: float = 1, center: bool = True, undirected: bool = False, distance=None
+): ...

stubs/networkx/networkx/generators/expanders.pyi

@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
 from typing_extensions import deprecated
 
-from networkx.classes.graph import Graph, _Node
+from networkx.classes.graph import Graph, _EdgeData, _Node, _NodeData
 from networkx.classes.multigraph import MultiGraph
 from networkx.utils.backends import _dispatchable
 
@@ -31,6 +31,6 @@ def maybe_regular_expander_graph(n: int, d: int, *, create_using=None, max_tries
 )
 def maybe_regular_expander(n, d, *, create_using=None, max_tries: int = 100, seed=None): ...
 @_dispatchable
-def is_regular_expander(G: Graph[_Node], *, epsilon: float = 0) -> bool: ...
+def is_regular_expander(G: Graph[_Node, _NodeData, _EdgeData], *, epsilon: float = 0) -> bool: ...
 @_dispatchable
 def random_regular_expander_graph(n: int, d: int, *, epsilon=0, create_using=None, max_tries=100, seed=None): ...

stubs/networkx/networkx/generators/geometric.pyi

@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
 from collections.abc import Callable, Iterable
 
-from networkx.classes.graph import Graph, _Node
+from networkx.classes.graph import Graph, _EdgeData, _Node, _NodeData
 from networkx.utils.backends import _dispatchable
 
 __all__ = [
@@ -16,7 +16,7 @@ __all__ = [
 ]
 
 @_dispatchable
-def geometric_edges(G: Graph[_Node], radius: float, p: float = 2) -> list[Incomplete]: ...
+def geometric_edges(G: Graph[_Node, _NodeData, _EdgeData], radius: float, p: float = 2) -> list[Incomplete]: ...
 @_dispatchable
 def random_geometric_graph(
     n: int | Iterable[Incomplete],

stubs/networkx/networkx/generators/line.pyi

@@ -1,11 +1,13 @@
 from _typeshed import Incomplete
 
-from networkx.classes.graph import Graph, _Node
+from networkx.classes.graph import Graph, _EdgeData, _Node, _NodeData
 from networkx.utils.backends import _dispatchable
 
 __all__ = ["line_graph", "inverse_line_graph"]
 
 @_dispatchable
-def line_graph(G: Graph[_Node], create_using: Graph[Incomplete] | type[Graph[Incomplete]] | None = None) -> Graph[Incomplete]: ...
+def line_graph(
+    G: Graph[_Node, _NodeData, _EdgeData], create_using: Graph[Incomplete] | type[Graph[Incomplete]] | None = None
+) -> Graph[Incomplete]: ...
 @_dispatchable
-def inverse_line_graph(G: Graph[_Node]) -> Graph[Incomplete]: ...
+def inverse_line_graph(G: Graph[_Node, _NodeData, _EdgeData]) -> Graph[Incomplete]: ...

stubs/networkx/networkx/generators/mycielski.pyi

@@ -1,11 +1,11 @@
 from _typeshed import Incomplete
 
-from networkx.classes.graph import Graph, _Node
+from networkx.classes.graph import Graph, _EdgeData, _Node, _NodeData
 from networkx.utils.backends import _dispatchable
 
 __all__ = ["mycielskian", "mycielski_graph"]
 
 @_dispatchable
-def mycielskian(G: Graph[_Node], iterations: int = 1) -> Graph[Incomplete]: ...
+def mycielskian(G: Graph[_Node, _NodeData, _EdgeData], iterations: int = 1) -> Graph[Incomplete]: ...
 @_dispatchable
 def mycielski_graph(n: int) -> Graph[Incomplete]: ...