Follow-up to graph layout cleanup: improvements for tests and typing (#3728)

* suggestions from review on #3434

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: behackl

manim/mobject/graph.py

@@ -9,11 +9,17 @@
 
 import itertools as it
 from copy import copy
-from typing import Any, Hashable, Iterable, Literal, Protocol, Union, cast
+from typing import TYPE_CHECKING, Any, Hashable, Iterable, Literal, Protocol, cast
 
 import networkx as nx
 import numpy as np
-from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    from typing_extensions import TypeAlias
+
+    from manim.typing import Point3D
+
+    NxGraph: TypeAlias = nx.classes.graph.Graph | nx.classes.digraph.DiGraph
 
 from manim.animation.composition import AnimationGroup
 from manim.animation.creation import Create, Uncreate
@@ -24,11 +30,8 @@
 from manim.mobject.opengl.opengl_mobject import OpenGLMobject
 from manim.mobject.text.tex_mobject import MathTex
 from manim.mobject.types.vectorized_mobject import VMobject
-from manim.typing import Point3D
 from manim.utils.color import BLACK
 
-NxGraph: TypeAlias = Union[nx.classes.graph.Graph, nx.classes.digraph.DiGraph]
-
 
 class LayoutFunction(Protocol):
     """A protocol for automatic layout functions that compute a layout for a graph to be used in :meth:`~.Graph.change_layout`.
@@ -53,8 +56,8 @@ def custom_layout(
                     graph: nx.Graph,
                     scale: float | tuple[float, float, float] = 2,
                     n: int | None = None,
-                    *args: tuple[Any, ...],
-                    **kwargs: dict[str, Any],
+                    *args: Any,
+                    **kwargs: Any,
                 ):
                     nodes = sorted(list(graph))
                     height = len(nodes) // n
@@ -256,8 +259,8 @@ def __call__(
         self,
         graph: NxGraph,
         scale: float | tuple[float, float, float] = 2,
-        *args: tuple[Any, ...],
-        **kwargs: dict[str, Any],
+        *args: Any,
+        **kwargs: Any,
     ) -> dict[Hashable, Point3D]:
         """Given a graph and a scale, return a dictionary of coordinates.
 
@@ -280,7 +283,7 @@ def _partite_layout(
     nx_graph: NxGraph,
     scale: float = 2,
     partitions: list[list[Hashable]] | None = None,
-    **kwargs: dict[str, Any],
+    **kwargs: Any,
 ) -> dict[Hashable, Point3D]:
     if partitions is None or len(partitions) == 0:
         raise ValueError(
@@ -302,7 +305,7 @@ def _partite_layout(
     return nx.layout.multipartite_layout(nx_graph, scale=scale, **kwargs)
 
 
-def _random_layout(nx_graph: NxGraph, scale: float = 2, **kwargs: dict[str, Any]):
+def _random_layout(nx_graph: NxGraph, scale: float = 2, **kwargs: Any):
     # the random layout places coordinates in [0, 1)
     # we need to rescale manually afterwards...
     auto_layout = nx.layout.random_layout(nx_graph, **kwargs)

tests/module/mobject/test_graph.py

@@ -116,9 +116,7 @@ def test_custom_graph_layout_dict():
 
 
 def test_graph_layouts():
-    for layout in (
-        layout for layout in _layouts if layout != "tree" and layout != "partite"
-    ):
+    for layout in (layout for layout in _layouts if layout not in ["tree", "partite"]):
         G = Graph([1, 2, 3], [(1, 2), (2, 3)], layout=layout)
         assert str(G) == "Undirected graph on 3 vertices and 2 edges"
 
@@ -164,9 +162,7 @@ def layout_func(graph, scale, offset):
 
 
 def test_graph_change_layout():
-    for layout in (
-        layout for layout in _layouts if layout != "tree" and layout != "partite"
-    ):
+    for layout in (layout for layout in _layouts if layout not in ["tree", "partite"]):
         G = Graph([1, 2, 3], [(1, 2), (2, 3)])
         G.change_layout(layout=layout)
         assert str(G) == "Undirected graph on 3 vertices and 2 edges"