aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorAleksandr <ivansduck@gmail.com>2022-06-01 13:47:52 +0300
committerAleksandr <ivansduck@gmail.com>2022-06-01 13:47:52 +0300
commit18e30b63b3358b58bcce35a95d876e459e602ef6 (patch)
tree96027ca9298f26a940ff85377f0949c999d80bb7
parent3087e92640f7488a160a24572ab02b0d8f55e1fc (diff)
downloadydb-18e30b63b3358b58bcce35a95d876e459e602ef6.tar.gz
DEVTOOLS-9500 [library/python/pytest] Generate call graph if --profile-pytest is specified
ref:b1b93724ca5ccf2a420b7dd063764bc8eed4f3e4
-rw-r--r--library/python/pytest/main.py25
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)