import tracing

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

2 files changed, 29 insertions, 2 deletions

diff --git a/library/python/testing/import_test/import_test.py b/library/python/testing/import_test/import_test.py
index 02331c18a7..7795c3c321 100644
--- a/library/python/testing/import_test/import_test.py
+++ b/library/python/testing/import_test/import_test.py
@@ -53,7 +53,7 @@ def check_imports(no_check=(), extra=(), skip_func=None, py_main=None):
     import_times = {}
 
     def norm(s):
-        return (s[:-9] if s.endswith('.__init__') else s)
+        return s[:-9] if s.endswith('.__init__') else s
 
     modules = sys.extra_modules | set(extra)
     modules = sorted(modules, key=norm)
@@ -87,7 +87,7 @@ def check_imports(no_check=(), extra=(), skip_func=None, py_main=None):
                 if module == '__main__':
                     importer.load_module('__main__', '__main__py')
                 elif module.endswith('.__init__'):
-                    __import__(module[:-len('.__init__')])
+                    __import__(module[: -len('.__init__')])
                 else:
                     __import__(module)
 
@@ -123,6 +123,29 @@ def main():
 
     skip_names = sys.argv[1:]
 
+    # SIGUSR2 is used by test_tool to teardown tests
+    if hasattr(signal, "SIGUSR2"):
+        # Dump python import tracing
+        import library.python.import_tracing.lib.regulator as regulator
+
+        # get the original handler to return control to it after dumping
+        signum = signal.SIGUSR2
+        orig_handler = signal.getsignal(signum)
+
+        if not hasattr(signal, 'raise_signal'):
+            # Only available for Python 3.8+
+            def raise_signal(signum):
+                os.kill(os.getpid(), signum)
+        else:
+            raise_signal = signal.raise_signal
+
+        def stop_tracing_handler(s, f):
+            regulator.disable(close_not_finished=True)
+            signal.signal(signal.SIGUSR2, orig_handler)
+            raise_signal(signum)
+
+        signal.signal(signal.SIGUSR2, stop_tracing_handler)
+
     try:
         import faulthandler
     except ImportError:
diff --git a/library/python/testing/import_test/ya.make b/library/python/testing/import_test/ya.make
index 3f95288f4d..504e3f3a57 100644
--- a/library/python/testing/import_test/ya.make
+++ b/library/python/testing/import_test/ya.make
@@ -2,4 +2,8 @@ PY23_LIBRARY()
 
 PY_SRCS(import_test.py)
 
+PEERDIR(
+    library/python/import_tracing/lib
+)
+
 END()