Add --gzip option to cli

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()

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()

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:

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

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__':