From 190ae5db53a76ac610d583394bc1e1fd0d236796 Mon Sep 17 00:00:00 2001 From: Nikhil Iyer Date: Thu, 2 Nov 2023 14:00:10 -0400 Subject: [PATCH] Fix None check order in _tree_layout (#3421) * Fix None check order in _tree_layout * add tests to test_graph.py * [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> --- manim/mobject/graph.py | 7 +++---- tests/module/mobject/test_graph.py | 14 ++++++++++++++ 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/manim/mobject/graph.py b/manim/mobject/graph.py index 45ada61936..12082d5ad4 100644 --- a/manim/mobject/graph.py +++ b/manim/mobject/graph.py @@ -112,13 +112,12 @@ def _tree_layout( vertex_spacing: tuple | None = None, orientation: str = "down", ): - children = {root_vertex: list(T.neighbors(root_vertex))} - - if not nx.is_tree(T): - raise ValueError("The tree layout must be used with trees") if root_vertex is None: raise ValueError("The tree layout requires the root_vertex parameter") + if not nx.is_tree(T): + raise ValueError("The tree layout must be used with trees") + children = {root_vertex: list(T.neighbors(root_vertex))} # The following code is SageMath's tree layout implementation, taken from # https://github.com/sagemath/sage/blob/cc60cfebc4576fed8b01f0fc487271bdee3cefed/src/sage/graphs/graph_plot.py#L1447 diff --git a/tests/module/mobject/test_graph.py b/tests/module/mobject/test_graph.py index 5ae3eda4fb..b05d6b21d9 100644 --- a/tests/module/mobject/test_graph.py +++ b/tests/module/mobject/test_graph.py @@ -1,5 +1,7 @@ from __future__ import annotations +import pytest + from manim import DiGraph, Graph, Scene, Text, tempconfig @@ -100,3 +102,15 @@ def test_custom_animation_mobject_list(): scene.play(G.animate.remove_vertices(2)) assert str(G) == "Undirected graph on 3 vertices and 0 edges" assert scene.mobjects == [G] + + +def test_tree_layout_no_root_error(): + with pytest.raises(ValueError) as excinfo: + G = Graph([1, 2, 3], [(1, 2), (2, 3)], layout="tree") + assert str(excinfo.value) == "The tree layout requires the root_vertex parameter" + + +def test_tree_layout_not_tree_error(): + with pytest.raises(ValueError) as excinfo: + G = Graph([1, 2, 3], [(1, 2), (2, 3), (3, 1)], layout="tree", root_vertex=1) + assert str(excinfo.value) == "The tree layout must be used with trees"