From 62ca889cb698e3b34fb00d2caee12c5fbdb43fda Mon Sep 17 00:00:00 2001
From: joshop <joshop4646@gmail.com>
Date: Wed, 9 Aug 2023 10:52:40 -0400
Subject: [PATCH] Clean up, CSV reader can be passed CSVs as strings

---
 pipit/trace.py               | 9 +++++++--
 pipit/util/test_generator.py | 7 +++----
 2 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/pipit/trace.py b/pipit/trace.py
index 4e320f12..f0f483fb 100644
--- a/pipit/trace.py
+++ b/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
diff --git a/pipit/util/test_generator.py b/pipit/util/test_generator.py
index be12c455..c6343256 100644
--- a/pipit/util/test_generator.py
+++ b/pipit/util/test_generator.py
@@ -352,11 +352,11 @@ 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
@@ -364,7 +364,6 @@ def generate_fake_test(
     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)