diff --git a/src/outliner/display.py b/src/outliner/display.py index 2c18daa..76c972a 100644 --- a/src/outliner/display.py +++ b/src/outliner/display.py @@ -36,6 +36,7 @@ def tree(self): connect_branch = ( lines + self.ansi_for_tree["CHILD_OBJECT_ANSI"] + + f"{position}. " + i + "\n" + self.ansi_for_tree["RESET"] @@ -44,7 +45,11 @@ def tree(self): connector_display = connector_root if position == 0 else connect_branch sys.stdout.write(connector_display) - sys.stdout.write(str(" " + self.ansi_for_tree["VERTICAL_LINE"] + "\n")) + + if position == (len(self.functions_flow) - 1): + return + + sys.stdout.write(str(" " + self.ansi_for_tree["VERTICAL_LINE"])) sys.stdout.flush() position += 1 diff --git a/tests/test_display.py b/tests/test_display.py index caafe28..e47ffd1 100644 --- a/tests/test_display.py +++ b/tests/test_display.py @@ -28,7 +28,7 @@ def setUp(self): @patch("sys.stdout", new_callable=StringIO) def test_display_tree_positive_input(self, stdout): - expected_result = " \x1b[38;5;208mtest2\x1b[0m\n │\n\n │──\x1b[36m__init__\n\x1b[0m │\n\n │──\x1b[36mfunc1\n\x1b[0m │\n\n │──\x1b[36mfunc2\n\x1b[0m │\n" + expected_result = " \x1b[38;5;208mtest2\x1b[0m\n │\n │──\x1b[36m1. __init__\n\x1b[0m │\n │──\x1b[36m2. func1\n\x1b[0m │\n │──\x1b[36m3. func2\n\x1b[0m" self.parser.tree() self.assertEqual(stdout.getvalue(), expected_result)