Fix tree printout for MultiIndex columns (#120)

LLNL · Apr 8, 2024 · fe68b12 · fe68b12
1 parent 79c009b
commit fe68b12
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/hatchet/external/console.py b/hatchet/external/console.py
@@ -327,7 +327,10 @@ def render_frame(self, node, dataframe, indent="", child_indent=""):
                 else:
                     metric_str += " [{}]".format(annotation_content)
 
-            node_name = dataframe.loc[df_index, self.name]
+            if isinstance(dataframe.columns, pd.MultiIndex):
+                node_name = dataframe.loc[df_index, (self.name, "")]
+            else:
+                node_name = dataframe.loc[df_index, self.name]
             if self.expand is False:
                 if len(node_name) > 39:
                     node_name = (