enhance visualization func type hint (#1443)

* enhance graphviz_draw type hint

* enhance mpl_draw type hint

* fix black lint

* fix black lint

Signed-off-by: ZhengYu, Xu <[email protected]>

* make Method and ImageType private to satisfy mypy.stubtest

* allow image_type and method to accept str

* remove type alias

* format

---------

Signed-off-by: ZhengYu, Xu <[email protected]>

Author: zen-xu

rustworkx/visualization/graphviz.pyi

@@ -10,17 +10,126 @@ import typing
 from rustworkx.rustworkx import PyGraph, PyDiGraph
 
 if typing.TYPE_CHECKING:
-    from PIL.Image import Image  # type: ignore
+    from PIL.Image import Image
 
 _S = typing.TypeVar("_S")
 _T = typing.TypeVar("_T")
 
+@typing.overload
 def graphviz_draw(
     graph: PyDiGraph[_S, _T] | PyGraph[_S, _T],
     node_attr_fn: typing.Callable[[_S], dict[str, str]] | None = ...,
     edge_attr_fn: typing.Callable[[_T], dict[str, str]] | None = ...,
     graph_attr: dict[str, str] | None = ...,
-    filename: str | None = ...,
+    filename: None = ...,
+    image_type: (
+        typing.Literal[
+            "canon",
+            "cmap",
+            "cmapx",
+            "cmapx_np",
+            "dia",
+            "dot",
+            "fig",
+            "gd",
+            "gd2",
+            "gif",
+            "hpgl",
+            "imap",
+            "imap_np",
+            "ismap",
+            "jpe",
+            "jpeg",
+            "jpg",
+            "mif",
+            "mp",
+            "pcl",
+            "pdf",
+            "pic",
+            "plain",
+            "plain-ext",
+            "png",
+            "ps",
+            "ps2",
+            "svg",
+            "svgz",
+            "vml",
+            "vmlzvrml",
+            "vtx",
+            "wbmp",
+            "xdor",
+            "xlib",
+        ]
+        | None
+    ) = ...,
+    method: typing.Literal["twopi", "neato", "circo", "fdp", "sfdp", "dot"] | None = ...,
+) -> Image: ...
+@typing.overload
+def graphviz_draw(
+    graph: PyDiGraph[_S, _T] | PyGraph[_S, _T],
+    node_attr_fn: typing.Callable[[_S], dict[str, str]] | None = ...,
+    edge_attr_fn: typing.Callable[[_T], dict[str, str]] | None = ...,
+    graph_attr: dict[str, str] | None = ...,
+    filename: None = ...,
+    image_type: str | None = ...,
+    method: str | None = ...,
+) -> Image: ...
+@typing.overload
+def graphviz_draw(
+    graph: PyDiGraph[_S, _T] | PyGraph[_S, _T],
+    node_attr_fn: typing.Callable[[_S], dict[str, str]] | None = ...,
+    edge_attr_fn: typing.Callable[[_T], dict[str, str]] | None = ...,
+    graph_attr: dict[str, str] | None = ...,
+    filename: str = ...,
+    image_type: (
+        typing.Literal[
+            "canon",
+            "cmap",
+            "cmapx",
+            "cmapx_np",
+            "dia",
+            "dot",
+            "fig",
+            "gd",
+            "gd2",
+            "gif",
+            "hpgl",
+            "imap",
+            "imap_np",
+            "ismap",
+            "jpe",
+            "jpeg",
+            "jpg",
+            "mif",
+            "mp",
+            "pcl",
+            "pdf",
+            "pic",
+            "plain",
+            "plain-ext",
+            "png",
+            "ps",
+            "ps2",
+            "svg",
+            "svgz",
+            "vml",
+            "vmlzvrml",
+            "vtx",
+            "wbmp",
+            "xdor",
+            "xlib",
+        ]
+        | None
+    ) = ...,
+    method: typing.Literal["twopi", "neato", "circo", "fdp", "sfdp", "dot"] | None = ...,
+) -> None: ...
+@typing.overload
+def graphviz_draw(
+    graph: PyDiGraph[_S, _T] | PyGraph[_S, _T],
+    node_attr_fn: typing.Callable[[_S], dict[str, str]] | None = ...,
+    edge_attr_fn: typing.Callable[[_T], dict[str, str]] | None = ...,
+    graph_attr: dict[str, str] | None = ...,
+    filename: str = ...,
     image_type: str | None = ...,
     method: str | None = ...,
-) -> Image | None: ...
+) -> None: ...

rustworkx/visualization/matplotlib.pyi

@@ -20,7 +20,7 @@ if typing.TYPE_CHECKING:
 _S = typing.TypeVar("_S")
 _T = typing.TypeVar("_T")
 
-class _DrawKwargs(typing.TypedDict, total=False):
+class _DrawKwargs(typing.TypedDict, typing.Generic[_S, _T], total=False):
     arrowstyle: str
     arrow_size: int
     node_list: list[int]
@@ -67,5 +67,5 @@ def mpl_draw(
     ax: Axes | None = ...,
     arrows: bool = ...,
     with_labels: bool = ...,
-    **kwds: typing_extensions.Unpack[_DrawKwargs],
+    **kwds: typing_extensions.Unpack[_DrawKwargs[_S, _T]],
 ) -> Figure | None: ...