diff --git a/example.py b/example.py index 3d1cec9f..e0798457 100644 --- a/example.py +++ b/example.py @@ -1,4 +1,3 @@ - def test_1(): return [a for a in range(1000000)] @@ -13,6 +12,10 @@ def main(): return test_1() + test_2() - -if __name__ == "__main__": +if __name__ == '__main__': + # When using `python -m vmprof` command main() + # When using vmprofi API. + #import vmprof + #with vmprof.profile("perf.data.gz", ["gzip", "-2"]): + # main() diff --git a/vmprof/__init__.py b/vmprof/__init__.py index eb9abf9f..80b9fd8c 100644 --- a/vmprof/__init__.py +++ b/vmprof/__init__.py @@ -1,5 +1,7 @@ import os +import subprocess import sys +import contextlib from . import cli @@ -47,3 +49,30 @@ def enable_jitlog(fileno): be broken. """ _vmprof.enable_jitlog(fileno) + + +@contextlib.contextmanager +def profile(outfile, pipecmd=None, period=DEFAULT_PERIOD, memory=False): + """Utility context manager which calls vmprof.enable() and vmprof.disable(). + + outfile is path to output file. + + This function support compression via pipe command:: + + with vmprof.profile("out.prof.gz", pipecmd=["/usr/bin/gzip", "-4"]): + main() + """ + with open(outfile, 'wb') as of: + proc = None + fileno = of.fileno() + if pipecmd is not None: + proc = subprocess.Popen(pipecmd, bufsize=-1, stdin=subprocess.PIPE, stdout=of.fileno()) + fileno = proc.stdin.fileno() + enable(fileno, period, memory) + try: + yield + finally: + disable() + if proc: + proc.stdin.close() + proc.wait() diff --git a/vmprof/__main__.py b/vmprof/__main__.py index eda22052..b1d89ebe 100644 --- a/vmprof/__main__.py +++ b/vmprof/__main__.py @@ -1,5 +1,7 @@ +import os import runpy -import sys, os +import subprocess +import sys import tempfile import vmprof @@ -40,6 +42,7 @@ def upload_stats(stats, forest, args): def main(): args = vmprof.cli.parse_args(sys.argv[1:]) + proc = None if args.web: output_mode = OUTPUT_WEB @@ -51,14 +54,23 @@ def main(): if output_mode == OUTPUT_FILE: prof_file = args.output prof_name = prof_file.name + fileno = prof_file.fileno() + if args.gzip: + cmd = ['/usr/bin/gzip', "-" + str(args.gzip)] + proc = subprocess.Popen(cmd, bufsize=-1, + stdin=subprocess.PIPE, + stdout=prof_file.fileno(), + close_fds=True) + fileno = proc.stdin.fileno() else: prof_file = tempfile.NamedTemporaryFile(delete=False) prof_name = prof_file.name + fileno = prof_file.fileno() if args.jitlog: assert hasattr(vmprof, 'enable_jitlog'), "note: jitlog is only available on pypy" - vmprof.enable(prof_file.fileno(), args.period, args.mem) + vmprof.enable(fileno, args.period, args.mem) if args.jitlog: # note that this file descr is then handled by jitlog fd = os.open(prof_name + '.jitlog', os.O_WRONLY | os.O_TRUNC | os.O_CREAT) @@ -74,6 +86,9 @@ def main(): vmprof.disable() if args.jitlog and hasattr(vmprof, 'disable_jitlog'): vmprof.disable_jitlog(fd) + if proc: + proc.stdin.close() + proc.wait() prof_file.close() show_stats(prof_name, output_mode, args) if output_mode != OUTPUT_FILE: diff --git a/vmprof/cli.py b/vmprof/cli.py index b875245f..c1ae3200 100644 --- a/vmprof/cli.py +++ b/vmprof/cli.py @@ -53,6 +53,14 @@ def build_argparser(): action='store_true', help='Upload the jitlog to remote server (defaults to vmprof.com)', ) + parser.add_argument( + '--gzip', + nargs='?', + type=int, + const=4, + metavar='level', + help="Compress output file with gzip with specified level", + ) output_mode_args = parser.add_mutually_exclusive_group() output_mode_args.add_argument( '--web', @@ -62,8 +70,8 @@ def build_argparser(): output_mode_args.add_argument( '--output', '-o', metavar='file.prof', - type=argparse.FileType('w+b'), - help='Save profiling data to file' + type=argparse.FileType('wb'), + help='Save profiling data to file', ) return parser diff --git a/vmprof/show.py b/vmprof/show.py index f2591bb7..c0d1b06a 100644 --- a/vmprof/show.py +++ b/vmprof/show.py @@ -1,9 +1,10 @@ from __future__ import absolute_import +import argparse import os import six +import sys import vmprof -import argparse class color(six.text_type): @@ -145,7 +146,10 @@ def main(): prune_level=args.prune_level, indent=args.indent) - pp.show(args.profile) + file = args.profile + if file == '-': + file = open(sys.stdin.fileno(), 'rb') + pp.show(file) if __name__ == '__main__':