From c500876b476059a39d30324a884e91d8ba1121e5 Mon Sep 17 00:00:00 2001
From: Christophe Bedard <fixed-term.christophe.bourquebedard@de.bosch.com>
Date: Thu, 6 Jun 2019 09:18:41 +0200
Subject: [PATCH] Add basic conversion entrypoint

---
 conversion/ctf.py | 61 +++++++++++++++++++++++++++++++++++++++++++++--
 convert.py        | 16 ++++++++++---
 2 files changed, 72 insertions(+), 5 deletions(-)

diff --git a/conversion/ctf.py b/conversion/ctf.py
index e5b233a..3ddba4e 100644
--- a/conversion/ctf.py
+++ b/conversion/ctf.py
@@ -1,3 +1,60 @@
+# CTF to pickle conversion
 
-def ctf_to_pickle():
-    print('ctf_to_pickle')
+import babeltrace
+from pickle import Pickler
+import time
+
+# TODO
+_IGNORED_FIELDS = []
+_DISCARD = "events_discarded"
+
+def ctf_to_pickle(trace_directory, target):
+    """
+    Load CTF trace and convert to a pickle file
+    :param trace_directory (str): the main/top trace directory
+    :param target (Pickler): the target pickle file to write to
+    """
+    # add traces
+    tc = babeltrace.TraceCollection()
+    print(f'Importing {trace_directory}')
+    tc.add_traces_recursive(trace_directory, 'ctf')
+
+    count = 0
+    count_written = 0
+    # count_pid_matched = 0
+    # traced = set()
+
+    start_time = time.time()
+
+    # PID_KEYS = ['vpid', 'pid']
+    for event in tc.events:
+        count += 1
+        # pid = None
+        # for key in PID_KEYS:
+        #     if key in event.keys():
+        #         pid = event[key]
+        #         break
+
+        # Write all for now
+        pod = _ctf_event_to_pod(event)
+        print(f'dumping pod: {str(pod)}')
+        target.dump(pod)
+        count_written += 1
+
+    print(f'{count_written} events in {time.time() - start_time}')
+
+
+def _ctf_event_to_pod(ctf_event):
+    """
+    Convert name, timestamp, and all other keys except those in IGNORED_FIELDS into a dictionary.
+    :param ctf_element: The element to convert
+    :type ctf_element: babeltrace.Element
+    :return:
+    :return type: dict
+    """
+    pod = {'_name': ctf_event.name, '_timestamp': ctf_event.timestamp}
+    if hasattr(ctf_event, _DISCARD) and ctf_event[_DISCARD] > 0:
+        print(ctf_event[_DISCARD])
+    for key in [key for key in ctf_event.keys() if key not in _IGNORED_FIELDS]:
+        pod[key] = ctf_event[key]
+    return pod
diff --git a/convert.py b/convert.py
index 3ab6144..3a41428 100644
--- a/convert.py
+++ b/convert.py
@@ -2,8 +2,18 @@
 # Entrypoint/script to convert CTF trace data to a pickle file
 # TODO
 
+import sys
+from pickle import Pickler
 from tracetools_analysis.conversion.ctf import *
 
-def main():
-    print('main called')
-    ctf_to_pickle()
+def main(argv=sys.argv):
+    if len(argv) != 3:
+        print("usage: /trace/directory pickle_target_file")
+        exit(1)
+
+    trace_directory = sys.argv[1]
+    pickle_target_file = sys.argv[2]
+
+    with open(pickle_target_file, 'wb') as f:
+        p = Pickler(f, protocol=4)
+        ctf_to_pickle(trace_directory, p)