import tracing

* add argument to parser
* add out_path as fn parameter
* set necessary env variables for import tracing

1 files changed, 52 insertions, 0 deletions

diff --git a/library/python/import_tracing/lib/import_tracer.py b/library/python/import_tracing/lib/import_tracer.py
new file mode 100644
index 0000000000..39dfc11bc9
--- /dev/null
+++ b/library/python/import_tracing/lib/import_tracer.py
@@ -0,0 +1,52 @@
+import threading
+import time
+import collections
+import library.python.import_tracing.lib.event as events
+import library.python.import_tracing.lib.constants as constants
+
+
+class ImportTracer:
+    def __init__(self):
+        self.events = collections.OrderedDict()
+        self.start_time = time.time()
+
+    def start_event(self, modname, filename, tid=None):
+        tid = tid if tid is not None else threading.current_thread().ident
+        time_from_start = self._get_current_time_from_start()
+
+        event_key = (modname, tid)
+        new_event = events.Event(
+            modname=modname,
+            filename=filename,
+            tid=tid,
+            start_time=time_from_start,
+            end_time=None,
+        )
+
+        self.events[event_key] = new_event
+
+    def finish_event(self, modname, filename, tid=None):
+        tid = tid if tid is not None else threading.current_thread().ident
+        event_key = (modname, tid)
+        event = self.events[event_key]
+
+        end_time = self._get_current_time_from_start()
+        event.end_time = end_time
+
+    def get_events(self, close_not_finished=False):
+        end_time = self._get_current_time_from_start()
+
+        for event in self.events.values():
+            if close_not_finished and event.end_time is None:
+                yield events.Event(
+                    modname=event.modname,
+                    filename=event.filename,
+                    tid=event.tid,
+                    start_time=event.start_time,
+                    end_time=end_time,
+                )
+            else:
+                yield event
+
+    def _get_current_time_from_start(self):
+        return (time.time() - self.start_time) * constants.MCS_IN_SEC