add render_header parameter to tree() (#66)

- toggle tree header on/off
- update tests

hatchet/external/console.py

@@ -44,7 +44,12 @@ def __init__(self, unicode=False, color=False):
         self.visited = []
 
     def render(self, roots, dataframe, **kwargs):
-        result = self.render_preamble()
+        self.render_header = kwargs["render_header"]
+
+        if self.render_header:
+            result = self.render_preamble()
+        else:
+            result = ""
 
         if roots is None:
             result += "The graph is empty.\n\n"

hatchet/graphframe.py

@@ -897,6 +897,7 @@ def tree(
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     ):
         """Format this graphframe as a tree and return the resulting string."""
         color = sys.stdout.isatty()
@@ -934,6 +935,7 @@ def tree(
             highlight_name=highlight_name,
             colormap=colormap,
             invert_colormap=invert_colormap,
+            render_header=render_header,
         )
 
     def to_dot(self, metric=None, name="name", rank=0, thread=0, threshold=0.0):

hatchet/tests/caliper.py

@@ -176,6 +176,7 @@ def test_tree(lulesh_caliper_json):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "121489.000 main" in output
     assert "663.000 LagrangeElements" in output
@@ -195,6 +196,7 @@ def test_tree(lulesh_caliper_json):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "662712.000 EvalEOSForElems" in output
     assert "2895319.000 LagrangeNodal" in output

hatchet/tests/cprofile.py

@@ -44,6 +44,7 @@ def test_tree(hatchet_cycle_pstats):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "g pstats_reader_test.py" in output
     assert "<method 'disable' ...Profiler' objects> ~" in output
@@ -62,6 +63,7 @@ def test_tree(hatchet_cycle_pstats):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "f pstats_reader_test.py" in output
     assert re.match("(.|\n)*recursive(.|\n)*recursive", output)

hatchet/tests/graphframe.py

@@ -20,6 +20,7 @@
 from hatchet.graph import Graph
 from hatchet.node import Node
 from hatchet.external.console import ConsoleRenderer
+from hatchet.version import __version__
 
 
 def test_copy(mock_graph_literal):
@@ -721,10 +722,12 @@ def test_tree(mock_graph_literal):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "0.000 foo" in output
     assert "10.000 waldo" in output
     assert "15.000 garply" in output
+    assert "v" + __version__ in output
 
     output = ConsoleRenderer(unicode=True, color=False).render(
         gf.graph.roots,
@@ -740,9 +743,11 @@ def test_tree(mock_graph_literal):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=False,
     )
     assert "55.000 waldo" in output
     assert "15.000 garply" in output
+    assert "v" + __version__ not in output
 
 
 def test_to_dot(mock_graph_literal):
@@ -798,6 +803,7 @@ def test_sub_decorator(small_mock1, small_mock2, small_mock3):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "0.000 C" in output
     assert u"nan D ▶" in output
@@ -827,6 +833,7 @@ def test_sub_decorator(small_mock1, small_mock2, small_mock3):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "15.000 A" in output
     assert u"5.000 C ◀" in output
@@ -861,6 +868,7 @@ def test_div_decorator(small_mock1, small_mock2):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "1.000 C" in output
     assert "inf B" in output

hatchet/tests/hpctoolkit.py

@@ -115,6 +115,7 @@ def test_tree(calc_pi_hpct_db):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "0.000 <program root> <unknown file>" in output
     assert (
@@ -136,6 +137,7 @@ def test_tree(calc_pi_hpct_db):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "17989.000 interp.c:0 interp.c" in output
     assert (

hatchet/tests/pyinstrument.py

@@ -91,6 +91,7 @@ def test_tree(hatchet_pyinstrument_json):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "0.000 <module> examples.py" in output
     assert "0.025 read hatchet/readers/caliper_reader.py" in output
@@ -109,6 +110,7 @@ def test_tree(hatchet_pyinstrument_json):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "0.478 <module> examples.py" in output
     assert "0.063 from_caliper_json hatchet/graphframe.py" in output

hatchet/tests/tau.py

@@ -43,6 +43,7 @@ def test_tree(tau_profile_dir):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "449.000 .TAU application" in output
     assert "4458.000 MPI_Finalize()" in output
@@ -63,6 +64,7 @@ def test_tree(tau_profile_dir):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
     assert "419.000 .TAU application" in output
     assert "4894.000 MPI_Finalize()" in output

hatchet/tests/timemory_test.py

@@ -58,6 +58,7 @@ def test_tree(timemory_json_data):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
 
     print(output)
@@ -76,6 +77,7 @@ def test_tree(timemory_json_data):
         highlight_name=False,
         colormap="RdYlGn",
         invert_colormap=False,
+        render_header=True,
     )
 
     print(output)