diff --git a/manim/mobject/graph.py b/manim/mobject/graph.py index 7eaee412b4..5c760b9632 100644 --- a/manim/mobject/graph.py +++ b/manim/mobject/graph.py @@ -670,6 +670,8 @@ def _populate_edge_dict( raise NotImplementedError("To be implemented in concrete subclasses") def __getitem__(self: Graph, v: Hashable) -> Mobject: + if isinstance(v, tuple) and len(v) == 2: + return self.edges[v] return self.vertices[v] def _create_vertex( diff --git a/tests/module/mobject/test_graph.py b/tests/module/mobject/test_graph.py index 16c32da88e..fb052119df 100644 --- a/tests/module/mobject/test_graph.py +++ b/tests/module/mobject/test_graph.py @@ -76,6 +76,19 @@ def test_graph_add_edges(): assert set(G._graph.edges()) == set(G.edges.keys()) +def test_graph_getitem(): + vertices = [1, 2, 3, 4] + edges = [(1, 2), (2, 3), (3, 4), (4, 1)] + G = Graph(vertices, edges) + # Vertex access + assert G[1] is G.vertices[1] + # Edge access via tuple key + assert G[(1, 2)] is G.edges[(1, 2)] + # DiGraph edge access + DG = DiGraph(vertices, edges) + assert DG[(1, 2)] is DG.edges[(1, 2)] + + def test_graph_remove_edges(): G = Graph([1, 2, 3, 4, 5], [(1, 2), (2, 3), (3, 4), (4, 5), (1, 5)]) removed_mobjects = G.remove_edges((1, 2))