DEVTOOLS-9500 [library/python/pytest] Generate call graph if --profile-pytest is specified

ref:b1b93724ca5ccf2a420b7dd063764bc8eed4f3e4

1 files changed, 25 insertions, 0 deletions

diff --git a/library/python/pytest/main.py b/library/python/pytest/main.py
index 6296bd6f0f..44046683eb 100644
--- a/library/python/pytest/main.py
+++ b/library/python/pytest/main.py
@@ -108,6 +108,31 @@ def main():
         profile.disable()
         ps = pstats.Stats(profile, stream=sys.stderr).sort_stats('cumulative')
         ps.print_stats()
+        if '--output-dir' in sys.argv:
+            output_dir = sys.argv[sys.argv.index('--output-dir') + 1]
+            prof_filename = os.path.join(output_dir, 'pytest.profile')
+            ps.dump_stats(prof_filename)
+
+            try:
+                import gprof2dot
+            except ImportError as e:
+                sys.stderr.write("Failed to generate call graph: {}\n".format(e))
+                gprof2dot = None
+
+            if gprof2dot:
+                import shlex
+
+                dot_filename = os.path.join(output_dir, 'pytest.profile.dot')
+                args = [
+                    prof_filename,
+                    '--format=pstats',
+                    '--output={}'.format(dot_filename),
+                ]
+                if 'PYTEST_GPROF2DOT_ARGS' in os.environ:
+                    x = os.environ['PYTEST_GPROF2DOT_ARGS']
+                    args.extend(shlex.split(x))
+
+                gprof2dot.main(argv=args)
 
     sys.exit(rc)