Merge pull request #255 from d-ryzhikov/async-profile

Add support for coroutine profiling

Author: fabianp

.gitignore

@@ -11,3 +11,4 @@ mprofile_*.dat
 
 # virtual environment
 venv/
+.python-version

.travis.yml

@@ -1,6 +1,5 @@
 language: python
 python:
-  - "2.7"
   - "3.4"
   - "3.5"
   - "3.6"

Makefile

@@ -1,6 +1,6 @@
 PYTHON ?= python
 
-.PHONY: test
+.PHONY: test develop
 
 test:
 	$(PYTHON) -m memory_profiler test/test_func.py
@@ -18,6 +18,7 @@ test:
 	$(PYTHON) test/test_exception.py
 	$(PYTHON) test/test_exit_code.py
 	$(PYTHON) test/test_mprof.py
+	$(PYTHON) test/test_async.py
 
 develop:
 	pip install -e .

memory_profiler.py

@@ -7,7 +7,10 @@
 
 _CMD_USAGE = "python -m memory_profiler script_file.py"
 
-from functools import wraps
+from asyncio import coroutine, iscoroutinefunction
+from contextlib import contextmanager
+from functools import partial, wraps
+import builtins
 import inspect
 import linecache
 import logging
@@ -18,7 +21,6 @@
 import time
 import traceback
 import warnings
-import contextlib
 
 if sys.platform == "win32":
     # any value except signal.CTRL_C_EVENT and signal.CTRL_BREAK_EVENT
@@ -43,17 +45,8 @@
     line_cell_magic = lambda func: func
     magics_class = lambda cls: cls
 
-PY2 = sys.version_info[0] == 2
-
 _TWO_20 = float(2 ** 20)
 
-if PY2:
-    import __builtin__ as builtins
-    to_str = lambda x: x
-    from future_builtins import filter
-else:
-    import builtins
-    to_str = lambda x: str(x)
 
 # .. get available packages ..
 try:
@@ -580,7 +573,7 @@ def f(*args, **kwds):
 
         return f
 
-    @contextlib.contextmanager
+    @contextmanager
     def call_on_stack(self, func, *args, **kwds):
         self.current_stack_level += 1
         self.stack[func].append(self.current_stack_level)
@@ -701,16 +694,27 @@ def add_function(self, func):
         else:
             self.code_map.add(code)
 
+    @contextmanager
+    def _count_ctxmgr(self):
+        self.enable_by_count()
+        try:
+            yield
+        finally:
+            self.disable_by_count()
+
     def wrap_function(self, func):
         """ Wrap a function to profile it.
         """
 
-        def f(*args, **kwds):
-            self.enable_by_count()
-            try:
-                return func(*args, **kwds)
-            finally:
-                self.disable_by_count()
+        if iscoroutinefunction(func):
+            @coroutine
+            def f(*args, **kwargs):
+                with self._count_ctxmgr():
+                    yield from func(*args, **kwargs)
+        else:
+            def f(*args, **kwds):
+                with self._count_ctxmgr():
+                    return func(*args, **kwds)
 
         return f
 
@@ -831,7 +835,7 @@ def show_results(prof, stream=None, precision=1):
                 inc = u''
                 occurences = u''
             tmp = template.format(lineno, total_mem, inc, occurences, all_lines[lineno - 1])
-            stream.write(to_str(tmp))
+            stream.write(tmp)
         stream.write(u'\n\n')
 
 
@@ -1121,12 +1125,23 @@ def profile(func=None, stream=None, precision=1, backend='psutil'):
         if not tracemalloc.is_tracing():
             tracemalloc.start()
     if func is not None:
-        @wraps(func)
-        def wrapper(*args, **kwargs):
-            prof = LineProfiler(backend=backend)
-            val = prof(func)(*args, **kwargs)
-            show_results(prof, stream=stream, precision=precision)
-            return val
+        get_prof = partial(LineProfiler, backend=backend)
+        show_results_bound = partial(
+            show_results, stream=stream, precision=precision
+        )
+        if iscoroutinefunction(func):
+            @coroutine
+            def wrapper(*args, **kwargs):
+                prof = get_prof()
+                val = yield from prof(func)(*args, **kwargs)
+                show_results_bound(prof)
+                return val
+        else:
+            def wrapper(*args, **kwargs):
+                prof = get_prof()
+                val = prof(func)(*args, **kwargs)
+                show_results_bound(prof)
+                return val
 
         return wrapper
     else:
@@ -1198,16 +1213,13 @@ def run_module_with_profiler(module, profiler, backend, passed_args=[]):
     ns = dict(_CLEAN_GLOBALS, profile=profiler)
     _backend = choose_backend(backend)
     sys.argv = [module] + passed_args
-    if PY2:
+    if _backend == 'tracemalloc' and has_tracemalloc:
+        tracemalloc.start()
+    try:
         run_module(module, run_name="__main__", init_globals=ns)
-    else:
-        if _backend == 'tracemalloc' and has_tracemalloc:
-            tracemalloc.start()
-        try:
-            run_module(module, run_name="__main__", init_globals=ns)
-        finally:
-            if has_tracemalloc and tracemalloc.is_tracing():
-                tracemalloc.stop()
+    finally:
+        if has_tracemalloc and tracemalloc.is_tracing():
+            tracemalloc.stop()
 
 
 class LogFile(object):

setup.py

@@ -1,22 +1,14 @@
-import os
-import io
 import re
 from setuptools import setup
 
 
 # https://packaging.python.org/guides/single-sourcing-package-version/
-def read(*names, **kwargs):
-    with io.open(
-        os.path.join(os.path.dirname(__file__), *names),
-        encoding=kwargs.get("encoding", "utf8")
-    ) as fp:
-        return fp.read()
-
-
-def find_version(*file_paths):
-    version_file = read(*file_paths)
-    version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
-                              version_file, re.M)
+def find_version(file_paths):
+    with open(file_paths) as f:
+        version_file = f.read()
+    version_match = re.search(
+        r"^__version__ = ['\"]([^'\"]*)['\"]", version_file, re.M
+    )
     if version_match:
         return version_match.group(1)
 
@@ -27,18 +19,14 @@ def find_version(*file_paths):
 Intended Audience :: Developers
 License :: OSI Approved :: BSD License
 Programming Language :: Python
-Programming Language :: Python :: 2
-Programming Language :: Python :: 2.6
-Programming Language :: Python :: 2.7
 Programming Language :: Python :: 3
-Programming Language :: Python :: 3.2
-Programming Language :: Python :: 3.3
 Topic :: Software Development
 Operating System :: POSIX
 Operating System :: Unix
 
 """
 
+
 setup(
     name='memory_profiler',
     description='A module for monitoring memory usage of a python program',
@@ -49,9 +37,10 @@ def find_version(*file_paths):
     url='https://github.com/pythonprofilers/memory_profiler',
     py_modules=['memory_profiler', 'mprof'],
     entry_points={
-        'console_scripts' : ['mprof = mprof:main'],
+        'console_scripts': ['mprof = mprof:main'],
     },
     install_requires=['psutil'],
+    python_requires='>=3.4',
     classifiers=[_f for _f in CLASSIFIERS.split('\n') if _f],
     license='BSD'
 )

test/test_async.py

@@ -0,0 +1,17 @@
+import asyncio
+
+from memory_profiler import profile
+
+
+@profile
+@asyncio.coroutine
+def my_func():
+    a = [1] * (10 ** 6)
+    b = [2] * (2 * 10 ** 7)
+    yield from asyncio.sleep(1e-2)
+    del b
+
+
+if __name__ == '__main__':
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(my_func())