networkx: complete the nx_latex module (#14581)

Author: hamdanal

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

@@ -0,0 +1,11 @@
+import networkx as nx
+
+# Test covariant dict type for `pos` in nx_latex functions
+G: "nx.Graph[int]" = nx.Graph([(1, 2), (2, 3), (3, 4)])
+nx.to_latex_raw(G, pos=nx.spring_layout(G, seed=42))  # OK: dict[node, ndarray]
+pos1: dict[int, tuple[int, int]] = {1: (1, 2), 2: (3, 4), 3: (5, 6), 4: (7, 8)}
+nx.to_latex_raw(G, pos=pos1)  # OK: dict[node, 2-tuple]
+pos2: dict[int, str] = {1: "(1, 2)", 2: "(3, 4)", 3: "(5, 6)", 4: "(7, 8)"}
+nx.to_latex_raw(G, pos=pos2)  # OK: dict[node, str]
+pos3: dict[int, int] = {1: 1, 2: 3, 3: 5, 4: 7}
+nx.to_latex_raw(G, pos=pos3)  # type: ignore # dict keys must be str or collection

stubs/networkx/networkx/drawing/nx_latex.pyi

@@ -1,36 +1,70 @@
+from _typeshed import StrPath, SupportsWrite
+from collections.abc import Collection
+from typing_extensions import TypeAlias, TypeVar
+
+from networkx.classes.graph import Graph, _Node
+
 __all__ = ["to_latex_raw", "to_latex", "write_latex"]
 
+# runtime requires a dict but it doesn't mutate it, we use a bounded typevar as
+# a values type to make type checkers treat the dict covariantely
+_PosT = TypeVar("_PosT", bound=Collection[float] | str)
+_Pos: TypeAlias = str | dict[_Node, _PosT]
+
 def to_latex_raw(
-    G,
-    pos: str = "pos",
+    G: Graph[_Node],
+    pos: _Pos[_Node, _PosT] = "pos",
     tikz_options: str = "",
     default_node_options: str = "",
-    node_options: str = "node_options",
-    node_label: str = "label",
+    node_options: str | dict[_Node, str] = "node_options",
+    node_label: str | dict[_Node, str] = "label",
     default_edge_options: str = "",
-    edge_options: str = "edge_options",
-    edge_label: str = "label",
-    edge_label_options: str = "edge_label_options",
-): ...
+    edge_options: str | dict[tuple[_Node, _Node], str] = "edge_options",
+    edge_label: str | dict[tuple[_Node, _Node], str] = "label",
+    edge_label_options: str | dict[tuple[_Node, _Node], str] = "edge_label_options",
+) -> str: ...
 def to_latex(
-    Gbunch,
-    pos: str = "pos",
+    Gbunch: Graph[_Node] | Collection[Graph[_Node]],
+    pos: _Pos[_Node, _PosT] | Collection[_Pos[_Node, _PosT]] = "pos",
+    tikz_options: str = "",
+    default_node_options: str = "",
+    node_options: str | dict[_Node, str] = "node_options",
+    node_label: str | dict[_Node, str] = "node_label",
+    default_edge_options: str = "",
+    edge_options: str | dict[tuple[_Node, _Node], str] = "edge_options",
+    edge_label: str | dict[tuple[_Node, _Node], str] = "edge_label",
+    edge_label_options: str | dict[tuple[_Node, _Node], str] = "edge_label_options",
+    caption: str = "",
+    latex_label: str = "",
+    sub_captions: Collection[str] | None = None,
+    sub_labels: Collection[str] | None = None,
+    n_rows: int = 1,
+    as_document: bool = True,
+    document_wrapper: str = ...,
+    figure_wrapper: str = ...,
+    subfigure_wrapper: str = ...,
+) -> str: ...
+def write_latex(
+    Gbunch: Graph[_Node] | Collection[Graph[_Node]],
+    path: StrPath | SupportsWrite[str],
+    *,
+    # **options passed to `to_latex`
+    pos: _Pos[_Node, _PosT] | Collection[_Pos[_Node, _PosT]] = "pos",
     tikz_options: str = "",
     default_node_options: str = "",
-    node_options: str = "node_options",
-    node_label: str = "node_label",
+    node_options: str | dict[_Node, str] = "node_options",
+    node_label: str | dict[_Node, str] = "node_label",
     default_edge_options: str = "",
-    edge_options: str = "edge_options",
-    edge_label: str = "edge_label",
-    edge_label_options: str = "edge_label_options",
+    edge_options: str | dict[tuple[_Node, _Node], str] = "edge_options",
+    edge_label: str | dict[tuple[_Node, _Node], str] = "edge_label",
+    edge_label_options: str | dict[tuple[_Node, _Node], str] = "edge_label_options",
     caption: str = "",
     latex_label: str = "",
-    sub_captions=None,
-    sub_labels=None,
+    sub_captions: Collection[str] | None = None,
+    sub_labels: Collection[str] | None = None,
     n_rows: int = 1,
     as_document: bool = True,
     document_wrapper: str = ...,
     figure_wrapper: str = ...,
     subfigure_wrapper: str = ...,
-): ...
-def write_latex(Gbunch, path, **options) -> None: ...
+) -> None: ...