Extract a function that selectively displays or prints DF (#6031)

Summary:
Pull Request resolved: #6031

The extracted function can be reused to selectively print or display other profiling events.

Reviewed By: Olivia-liu

Differential Revision: D63468910

fbshipit-source-id: 71e2271dcaf9d7871e0e4289348bb5298677f6a1

Author: Margarita Grinvald

devtools/inspector/_inspector.py

@@ -43,6 +43,7 @@
 from executorch.devtools.inspector._inspector_utils import (
     calculate_time_scale_factor,
     create_debug_handle_to_op_node_mapping,
+    display_or_print_df,
     EDGE_DIALECT_GRAPH_KEY,
     EXCLUDED_COLUMNS_WHEN_PRINTING,
     EXCLUDED_EVENTS_WHEN_PRINTING,
@@ -60,8 +61,6 @@
 )
 from executorch.exir import ExportedProgram
 
-from tabulate import tabulate
-
 
 log: logging.Logger = logging.getLogger(__name__)
 
@@ -1172,25 +1171,7 @@ def print_data_tabular(
             ]
         filtered_column_df.reset_index(drop=True, inplace=True)
 
-        try:
-            from IPython import get_ipython
-            from IPython.display import display
-
-            def style_text_size(val, size=12):
-                return f"font-size: {size}px"
-
-            if get_ipython() is not None:
-                styled_df = filtered_column_df.style.applymap(style_text_size)
-                display(styled_df)
-            else:
-                raise Exception(
-                    "Environment unable to support IPython. Fall back to print()."
-                )
-        except:
-            print(
-                tabulate(filtered_column_df, headers="keys", tablefmt="fancy_grid"),
-                file=file,
-            )
+        display_or_print_df(filtered_column_df, file)
 
     # TODO: write unit test
     def find_total_for_module(self, module_name: str) -> float:

devtools/inspector/_inspector_utils.py

@@ -7,11 +7,14 @@
 # pyre-unsafe
 
 import math
+import sys
 from enum import Enum
-from typing import Dict, List, Mapping, Optional, Tuple, TypeAlias, Union
+from typing import Dict, IO, List, Mapping, Optional, Tuple, TypeAlias, Union
 
 import executorch.devtools.etdump.schema_flatcc as flatcc
 
+import pandas as pd
+
 import torch
 
 from executorch.devtools.debug_format.base_schema import OperatorNode
@@ -30,6 +33,8 @@
 from executorch.devtools.etdump.serialize import deserialize_from_etdump_flatcc
 from executorch.devtools.etrecord import ETRecord
 
+from tabulate import tabulate
+
 FORWARD = "forward"
 EDGE_DIALECT_GRAPH_KEY = "edge_dialect_graph_module"
 
@@ -295,6 +300,28 @@ def gen_etdump_object(
     return deserialize_from_etdump_flatcc(etdump_data)
 
 
+def display_or_print_df(df: pd.DataFrame, file: IO[str] = sys.stdout):
+    try:
+        from IPython import get_ipython
+        from IPython.display import display
+
+        def style_text_size(val, size=12):
+            return f"font-size: {size}px"
+
+        if get_ipython() is not None:
+            styled_df = df.style.applymap(style_text_size)
+            display(styled_df)
+        else:
+            raise Exception(
+                "Environment unable to support IPython. Fall back to print()."
+            )
+    except:
+        print(
+            tabulate(df, headers="keys", tablefmt="fancy_grid"),
+            file=file,
+        )
+
+
 def plot_metric(result: List[float], metric_name: str):
     import matplotlib.pyplot as plt
     import numpy as np