Add coroutine support

Author: d-ryzhykau

_aio_34.py

@@ -17,14 +17,19 @@
 import warnings
 import contextlib
 
-from .common import PY2, HAS_TRACEMALLOC
+from .common import PY2, PY34, PY35, HAS_TRACEMALLOC
 from .utils import (
     show_results,
     choose_backend,
     get_memory as _get_memory,
     get_child_memory as _get_child_memory,
 )
-from .line_profiler import LineProfiler
+from .line_profiler import LineProfiler, get_profile_wrapper
+
+if PY34:
+    from ._aio_34 import get_profile_wrapper
+elif PY35:
+    from ._aio_35 import get_profile_wrapper
 
 
 if sys.platform == "win32":
@@ -758,14 +763,7 @@ 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
-
-        return wrapper
+        return get_profile_wrapper(func, precision, backend, stream)
     else:
         def inner_wrapper(f):
             return profile(f, stream=stream, precision=precision,

memory_profiler/__init__.py

@@ -0,0 +1,35 @@
+from asyncio import coroutine, iscoroutinefunction
+from functools import wraps
+
+from .line_profiler import (
+    get_profile_wrapper as default_profile_wrapper,
+    LineProfiler,
+)
+from .utils import show_results
+
+
+class CoroLineProfiler(LineProfiler):
+    def wrap_function(self, func):
+        if iscoroutinefunction(func):
+            @coroutine
+            def f(*args, **kwargs):
+                with self.count_ctxmgr():
+                    yield from func(*args, **kwargs)
+            return f
+        else:
+            return super().wrap_function(func)
+
+
+def get_profile_wrapper(func, precision, backend, stream):
+    if iscoroutinefunction(func):
+        @wraps(func)
+        @coroutine
+        def wrapper(*args, **kwargs):
+            prof = CoroLineProfiler(backend=backend)
+            val = yield from prof(func)(*args, **kwargs)
+            show_results(prof, stream=stream, precision=precision)
+            return val
+    else:
+        wrapper = default_profile_wrapper(func, precision, backend, stream)
+
+    return wrapper

memory_profiler/_aio_34.py

@@ -0,0 +1,33 @@
+from functools import wraps
+from inspect import iscoroutinefunction
+
+from .line_profiler import (
+    get_profile_wrapper as default_profile_wrapper,
+    LineProfiler,
+)
+from .utils import show_results
+
+
+class CoroLineProfiler(LineProfiler):
+    def wrap_function(self, func):
+        if iscoroutinefunction(func):
+            async def f(*args, **kwargs):
+                with self.count_ctxmgr():
+                    return await func(*args, **kwargs)
+            return f
+        else:
+            return super().wrap_function(func)
+
+
+def get_profile_wrapper(func, precision, backend, stream):
+    if iscoroutinefunction(func):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            prof = CoroLineProfiler(backend=backend)
+            val = await prof(func)(*args, **kwargs)
+            show_results(prof, stream=stream, precision=precision)
+            return val
+    else:
+        wrapper = default_profile_wrapper(func, precision, backend, stream)
+
+    return wrapper

memory_profiler/_aio_35.py

@@ -3,6 +3,9 @@
 
 PY2 = sys.version_info[0] == 2
 
+PY34 = (3, 4) < sys.version_info < (3, 5)
+PY35 = (3, 5) < sys.version_info
+
 try:
     import tracemalloc  # noqa
 except ImportError: