diff --git a/setup.py b/setup.py index 69bbd7b0..7b8f2a4d 100644 --- a/setup.py +++ b/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': [ diff --git a/tox.ini b/tox.ini index fa670710..e2db4a9e 100644 --- a/tox.ini +++ b/tox.ini @@ -1,7 +1,7 @@ [tox] -envlist = py27, py34, pypy +envlist = py27, py34, py35, pypy [testenv] deps=pytest -commands=py.test tests/ \ No newline at end of file +commands=py.test vmprof/test/ diff --git a/vmprof/__init__.py b/vmprof/__init__.py index 1bce7b76..718bc483 100644 --- a/vmprof/__init__.py +++ b/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 diff --git a/vmprof/reader.py b/vmprof/reader.py index aac1997a..2d39010a 100644 --- a/vmprof/reader.py +++ b/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 diff --git a/vmprof/test/test_run.py b/vmprof/test/test_run.py index 7a8b012d..03ccab34 100644 --- a/vmprof/test/test_run.py +++ b/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()