tree: Avoid extra trailing spaces on <ROOT> node

If trailing spaces are included, they are highlighted with an orange background when the line is selected. Signed-off-by: Matt Wozniski <[email protected]>
bloomberg · Jul 2, 2024 · 94fc8ed · 94fc8ed
1 parent f11a9e3
commit 94fc8ed
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/src/memray/reporters/tree.py b/src/memray/reporters/tree.py
@@ -308,11 +308,9 @@ def frame_text(self, node: Frame, *, allow_expand: bool) -> Text:
             code_position = (
                 f"{_filename_to_module_name(file)}:{lineno}" if lineno != 0 else file
             )
-            ret.append_text(
-                Text.from_markup(
-                    f"[bold]{function}[/bold]  [dim cyan]{code_position}[/dim cyan]"
-                )
-            )
+            ret.append_text(Text.from_markup(f"[bold]{function}[/]"))
+            if code_position:
+                ret.append_text(Text.from_markup(f"  [dim cyan]{code_position}[/]"))
         else:
             ret.append_text(Text("hidden"))
         return ret

diff --git a/tests/unit/test_tree_reporter.py b/tests/unit/test_tree_reporter.py
@@ -853,7 +853,7 @@ async def run_test():
 
         # THEN
         assert tree_to_dict(root) == TreeElement(
-            label="📂 1.000KB (100.00 %) <ROOT>  ",
+            label="📂 1.000KB (100.00 %) <ROOT>",
             children=[
                 TreeElement(
                     label="📂 1.000KB (100.00 %) grandparent  fun.py:4",
@@ -940,7 +940,7 @@ async def run_test():
 
         # THEN
         assert tree_to_dict(root) == TreeElement(
-            label="📂 10.000KB (100.00 %) <ROOT>  ",
+            label="📂 10.000KB (100.00 %) <ROOT>",
             children=[
                 TreeElement(
                     label="📂 6.000KB (60.00 %) grandparent2  fun2.py:4",
@@ -1045,7 +1045,7 @@ async def run_test():
 
         # THEN
         assert first_tree == TreeElement(
-            label="📂 11.000KB (100.00 %) <ROOT>  ",
+            label="📂 11.000KB (100.00 %) <ROOT>",
             children=[
                 TreeElement(
                     label="📂 10.000KB (90.91 %) grandparent2  fun2.py:4",
@@ -1078,7 +1078,7 @@ async def run_test():
             is_expanded=True,
         )
         assert second_tree == TreeElement(
-            label="📂 11.000KB (100.00 %) <ROOT>  ",
+            label="📂 11.000KB (100.00 %) <ROOT>",
             children=[
                 TreeElement(
                     label="📂 10.000KB (90.91 %) grandparent2  fun2.py:4",
@@ -1214,7 +1214,7 @@ async def run_test():
         # THEN
 
         assert tree_to_dict(root) == TreeElement(
-            label="📂 1.000KB (100.00 %) <ROOT>  ",
+            label="📂 1.000KB (100.00 %) <ROOT>",
             children=[
                 TreeElement(
                     label="📄 1.000KB (100.00 %) me  foo.py:12",
@@ -1278,7 +1278,7 @@ async def run_test():
 
         # THEN
         assert first_tree == TreeElement(
-            label="📂 11.000KB (100.00 %) <ROOT>  ",
+            label="📂 11.000KB (100.00 %) <ROOT>",
             children=[
                 TreeElement(
                     label="📂 10.000KB (90.91 %) grandparent2  fun2.py:4",
@@ -1305,7 +1305,7 @@ async def run_test():
             is_expanded=True,
         )
         assert second_tree == TreeElement(
-            label="📂 11.000KB (100.00 %) <ROOT>  ",
+            label="📂 11.000KB (100.00 %) <ROOT>",
             children=[
                 TreeElement(
                     label="📂 10.000KB (90.91 %) grandparent2  fun2.py:4",
@@ -1437,7 +1437,7 @@ async def run_test():
 
         # THEN
         assert tree_to_dict(root) == TreeElement(
-            label="📂 11.000KB (100.00 %) <ROOT>  ",
+            label="📂 11.000KB (100.00 %) <ROOT>",
             children=[
                 TreeElement(
                     label="📂 10.000KB (90.91 %) grandparent2  fun2.py:4",
@@ -1510,7 +1510,7 @@ async def run_test():
 
         root = async_run(run_test())
 
-        assert str(root.label) == "📂 1.000KB (100.00 %) <ROOT>  "
+        assert str(root.label) == "📂 1.000KB (100.00 %) <ROOT>"
         assert len(root.children) == 1
         current_node = root.children[0]
         for i in range(1, MAX_STACKS + 2):