From 9bc79226132198636886dc9aa1c72939f356570e Mon Sep 17 00:00:00 2001
From: "Yury A. Apollov" <apollovy@gmail.com>
Date: Fri, 29 Sep 2017 20:02:26 +0300
Subject: [PATCH] Add handling of asyncio coroutines.

---
 memory_profiler.py | 57 +++++++++++++++++++++++++++++++++-------------
 setup.py           |  9 ++++----
 2 files changed, 46 insertions(+), 20 deletions(-)

diff --git a/memory_profiler.py b/memory_profiler.py
index 9e74846..042e66c 100644
--- a/memory_profiler.py
+++ b/memory_profiler.py
@@ -1,6 +1,9 @@
 """Profile the memory usage of a Python program"""
 
+import asyncio
 # .. we'll use this to pass it to the child script ..
+import functools
+
 _CLEAN_GLOBALS = globals().copy()
 
 __version__ = '0.47'
@@ -16,7 +19,6 @@
 import inspect
 import subprocess
 import logging
-import traceback
 from signal import SIGKILL
 
 
@@ -654,13 +656,24 @@ 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 not asyncio.iscoroutinefunction(func):
+            def f(*args, **kwds):
+                self.enable_by_count()
+                try:
+                    return func(*args, **kwds)
+                finally:
+                    self.disable_by_count()
+        else:
+            @asyncio.coroutine
+            def f(*args, **kwds):
+                self.enable_by_count()
+                try:
+                    val = yield from func(*args, **kwds)
+                    return val
+                finally:
+                    self.disable_by_count()
+
+        functools.update_wrapper(f, func)
         return f
 
     def runctx(self, cmd, globals, locals):
@@ -1062,20 +1075,32 @@ def profile(func=None, stream=None, precision=1, backend='psutil'):
     if backend == 'tracemalloc' and has_tracemalloc:
         if not tracemalloc.is_tracing():
             tracemalloc.start()
+
     if func is not None:
-        def wrapper(*args, **kwargs):
-            prof = LineProfiler(backend=backend)
-            val = prof(func)(*args, **kwargs)
-            show_results(prof, stream=stream, precision=precision)
-            return val
+        if not asyncio.iscoroutinefunction(func):
+            def wrapper(*args, **kwargs):
+                prof = LineProfiler(backend=backend)
+                val = prof(func)(*args, **kwargs)
+                show_results(prof, stream=stream, precision=precision)
+
+                return val
+        else:
+            @asyncio.coroutine
+            def wrapper(*args, **kwargs):
+                prof = LineProfiler(backend=backend)
+                val = yield from prof(func)(*args, **kwargs)
+                show_results(prof, stream=stream, precision=precision)
+
+                return val
 
-        return wrapper
     else:
-        def inner_wrapper(f):
+        def wrapper(f):
             return profile(f, stream=stream, precision=precision,
                            backend=backend)
 
-        return inner_wrapper
+    functools.update_wrapper(wrapper, func)
+
+    return wrapper
 
 
 def choose_backend(new_backend=None):
diff --git a/setup.py b/setup.py
index bc119fd..3124124 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,7 @@
-import memory_profiler
+import sys
 from distutils.core import setup
-import setuptools
+
+import memory_profiler
 
 CLASSIFIERS = """\
 Development Status :: 5 - Production/Stable
@@ -31,6 +32,6 @@
     py_modules=['memory_profiler'],
     scripts=['mprof'],
     classifiers=[_f for _f in CLASSIFIERS.split('\n') if _f],
-    license='BSD'
-
+    license='BSD',
+    install_requires=['asyncio>=3.4.3'] if sys.version_info < (3, 4) else [],
 )