Clean up, CSV reader can be passed CSVs as strings

pipit/trace.py

@@ -5,7 +5,8 @@
 
 import numpy as np
 import pandas as pd
-import ast
+from ast import literal_eval
+from io import StringIO
 
 
 class Trace:
@@ -61,6 +62,10 @@ def from_nsight(filename):
 
     @staticmethod
     def from_csv(filename):
+        # detect if the input is a CSV as a string
+        if "," in filename:
+            # wrapping with StringIO allows pandas to read it
+            filename = StringIO(filename)
         events_dataframe = pd.read_csv(filename, skipinitialspace=True)
 
         # if timestamps are in seconds, convert them to nanoseconds
@@ -78,7 +83,7 @@ def from_csv(filename):
         if "Attributes" in events_dataframe.columns:
             # use literal_eval so we're not running a security risk
             events_dataframe["Attributes"] = events_dataframe["Attributes"].apply(
-                ast.literal_eval
+                literal_eval
             )
 
         # make certain columns categorical

pipit/util/test_generator.py

@@ -352,19 +352,18 @@ def emit_tree_data(trees):
     return data_csv, ground_csv
 
 
-def generate_fake_test(
+def generate_trace(
     num_events,
     num_processes,
     function_names=["foo", "bar", "baz", "quux", "grault", "garply", "waldo"],
-    num_mpi_events=0,
+    num_mpi_pairs=0,
 ):
     """
     Top level test generation function. Generates test and ground truth datasets with a
     minimum of num_events Enter/Leave events per process, of which there are
     num_processes. Optionally, MPI events can be added.
     """
     seed_tree = gen_fake_tree(num_events // 2, function_names)
-    print(num_events // 2, seed_tree.total_nodes)
     forest = gen_forest(seed_tree, num_processes)
-    add_fake_mpi_events(forest, num_mpi_events)
+    add_fake_mpi_events(forest, num_mpi_pairs)
     return emit_tree_data(forest)