Merge pull request #90 from jonashaag/gzip-subproc-2

Gzip subproc. I'll release it in a pre release version shortly.

Author: planrich

setup.py

@@ -1,7 +1,10 @@
 from setuptools import setup, find_packages, Extension
 import os, sys
 
-if '__pypy__' in sys.builtin_module_names:
+PY3  = sys.version_info[0] >= 3
+IS_PYPY = '__pypy__' in sys.builtin_module_names
+
+if IS_PYPY:
     ext_modules = [] # built-in
 else:
     if sys.platform != 'win32':
@@ -21,6 +24,11 @@
                             extra_compile_args=extra_compile_args,
                             libraries=[])]
 
+if PY3:
+    extra_install_requires = []
+else:
+    extra_install_requires = ["backports.shutil_which"]
+
 setup(
     name='vmprof',
     author='vmprof team',
@@ -33,7 +41,7 @@
     install_requires=[
         'requests',
         'six',
-    ],
+    ] + extra_install_requires,
     tests_require=['pytest'],
     entry_points = {
         'console_scripts': [

tox.ini

@@ -1,7 +1,7 @@
 [tox]
-envlist = py27, py34, pypy
+envlist = py27, py34, py35, pypy
 
 
 [testenv]
 deps=pytest
-commands=py.test tests/
+commands=py.test vmprof/test/

vmprof/__init__.py

@@ -1,5 +1,10 @@
 import os
 import sys
+import subprocess
+try:
+    from shutil import which
+except ImportError:
+    from backports.shutil_which import which
 
 from . import cli
 
@@ -24,20 +29,50 @@
     def enable(fileno, period=DEFAULT_PERIOD, memory=False, lines=False):
         if not isinstance(period, float):
             raise ValueError("You need to pass a float as an argument")
-        _vmprof.enable(fileno, period, memory, lines)
-
-    def disable():
-        _vmprof.disable()
+        gz_fileno = _gzip_start(fileno)
+        _vmprof.enable(gz_fileno, period, memory, lines)
 else:
     def enable(fileno, period=DEFAULT_PERIOD, memory=False, lines=False, warn=True):
         if not isinstance(period, float):
             raise ValueError("You need to pass a float as an argument")
         if warn and sys.pypy_version_info[:3] < (4, 1, 0):
-            print("PyPy <4.1 have various kinds of bugs, pass warn=False if you know what you're doing\n")
             raise Exception("PyPy <4.1 have various kinds of bugs, pass warn=False if you know what you're doing")
+        if warn and memory:
+            print("Memory profiling is currently unsupported for PyPy. Running without memory statistics.")
         if warn and lines:
             print('Line profiling is currently unsupported for PyPy. Running without lines statistics.\n')
-        _vmprof.enable(fileno, period)
+        gz_fileno = _gzip_start(fileno)
+        _vmprof.enable(gz_fileno, period)
 
-    def disable():
+def disable():
+    try:
         _vmprof.disable()
+        _gzip_finish()
+    except IOError as e:
+        raise Exception("Error while writing profile: " + str(e))
+
+
+_gzip_proc = None
+
+def _gzip_start(fileno):
+    """Spawn a gzip subprocess that writes compressed profile data to `fileno`.
+
+    Return the subprocess' input fileno.
+    """
+    # Prefer system gzip and fall back to Python's gzip module
+    if which("gzip"):
+        gzip_cmd = ["gzip", "-", "-4"]
+    else:
+        gzip_cmd = ["python", "-m", "gzip"]
+    global _gzip_proc
+    _gzip_proc = subprocess.Popen(gzip_cmd, stdin=subprocess.PIPE,
+                                  stdout=fileno, bufsize=-1,
+                                  close_fds=(sys.platform != "win32"))
+    return _gzip_proc.stdin.fileno()
+
+def _gzip_finish():
+    global _gzip_proc
+    if _gzip_proc is not None:
+        _gzip_proc.stdin.close()
+        _gzip_proc.wait()
+        _gzip_proc = None

vmprof/reader.py

@@ -1,15 +1,18 @@
 from __future__ import print_function
 import re
+import os
 import struct
 import subprocess
 import sys
 from six.moves import xrange
+import io
+import gzip
 
-PY3 = sys.version_info[0] >= 3
+from vmprof.binary import read_word, read_string, read_words
 
+PY3  = sys.version_info[0] >= 3
 WORD_SIZE = struct.calcsize('L')
 
-from vmprof.binary import read_word, read_string, read_words
 
 
 def read_trace(fileobj, depth, version, profile_lines=False):
@@ -30,6 +33,7 @@ def read_trace(fileobj, depth, version, profile_lines=False):
                 trace[i] = -trace[i]
         return trace
 
+
 MARKER_STACKTRACE = b'\x01'
 MARKER_VIRTUAL_IP = b'\x02'
 MARKER_TRAILER = b'\x03'
@@ -53,12 +57,14 @@ def read_trace(fileobj, depth, version, profile_lines=False):
 VMPROF_GC_TAG = 5
 VMPROF_ASSEMBLER_TAG = 6
 
+
 class AssemblerCode(int):
     pass
 
 class JittedCode(int):
     pass
 
+
 def wrap_kind(kind, pc):
     if kind == VMPROF_ASSEMBLER_TAG:
         return AssemblerCode(pc)
@@ -67,6 +73,15 @@ def wrap_kind(kind, pc):
     assert kind == VMPROF_CODE_TAG
     return pc
 
+
+def gunzip(fileobj):
+    is_gzipped = fileobj.read(2) == b'\037\213'
+    fileobj.seek(-2, os.SEEK_CUR)
+    if is_gzipped:
+        fileobj = io.BufferedReader(gzip.GzipFile(fileobj=fileobj))
+    return fileobj
+
+
 class BufferTooSmallError(Exception):
     def get_buf(self):
         return b"".join(self.args[0])
@@ -174,6 +189,7 @@ def read_one_marker(fileobj, status, buffer_so_far=None):
     return False
 
 def read_prof_bit_by_bit(fileobj):
+    fileobj = gunzip(fileobj)
     # note that we don't want to use all of this on normal files, since it'll
     # cost us quite a bit in memory and performance and parsing 200M files in
     # CPython is slow (pypy does better, use pypy)
@@ -193,7 +209,9 @@ def read_prof_bit_by_bit(fileobj):
             buf = e.get_buf()
     return status.period, status.profiles, status.virtual_ips, status.interp_name
 
-def read_prof(fileobj, virtual_ips_only=False): #
+def read_prof(fileobj, virtual_ips_only=False):
+    fileobj = gunzip(fileobj)
+
     assert read_word(fileobj) == 0 # header count
     assert read_word(fileobj) == 3 # header size
     assert read_word(fileobj) == 0

vmprof/test/test_run.py

@@ -1,10 +1,10 @@
 
 """ Test the actual run
 """
-
 import py
 import sys
 import tempfile
+import gzip
 
 import six
 
@@ -22,7 +22,7 @@
     COUNT = 100000
 else:
     COUNT = 10000
-    
+
 def function_foo():
     for k in range(1000):
         l = [a for a in xrange(COUNT)]
@@ -45,7 +45,7 @@ def test_basic():
     function_foo()
     vmprof.disable()
     tmpfile.close()
-    assert b"function_foo" in open(tmpfile.name, 'rb').read()
+    assert b"function_foo" in gzip.GzipFile(tmpfile.name).read()
 
 def test_read_bit_by_bit():
     tmpfile = tempfile.NamedTemporaryFile(delete=False)
@@ -155,6 +155,15 @@ def function_bar():
 
     s = prof.get_stats()
 
+def test_gzip_problem():
+    tmpfile = tempfile.NamedTemporaryFile(delete=False)
+    vmprof.enable(tmpfile.fileno())
+    vmprof._gzip_proc.kill()
+    function_foo()
+    with py.test.raises(Exception) as exc_info:
+        vmprof.disable()
+        assert "Error while writing profile" in str(exc_info)
+    tmpfile.close()
 
 def test_line_profiling():
     tmpfile = tempfile.NamedTemporaryFile(delete=False)
@@ -175,6 +184,5 @@ def walk(tree):
         walk(stats.get_tree())
 
 
-
 if __name__ == '__main__':
     test_line_profiling()