Merge pull request #75 from pfreixes/master

Get support for new 3.5 Python coroutines

line_profiler.py

@@ -27,6 +27,7 @@
 # Python 2/3 compatibility utils
 # ===========================================================
 PY3 = sys.version_info[0] == 3
+PY35 = PY3 and sys.version_info[1] >= 5
 
 # exec (from https://bitbucket.org/gutworth/six/):
 if PY3:
@@ -46,6 +47,14 @@ def exec_(_code_, _globs_=None, _locs_=None):
             _locs_ = _globs_
         exec("""exec _code_ in _globs_, _locs_""")
 
+if PY35:
+    import inspect
+    def is_coroutine(f):
+        return inspect.iscoroutinefunction(f)
+else:
+    def is_coroutine(f):
+        return False
+
 # ============================================================
 
 CO_GENERATOR = 0x0020
@@ -65,7 +74,9 @@ def __call__(self, func):
         it on function exit.
         """
         self.add_function(func)
-        if is_generator(func):
+        if is_coroutine(func):
+            wrapper = self.wrap_coroutine(func)
+        elif is_generator(func):
             wrapper = self.wrap_generator(func)
         else:
             wrapper = self.wrap_function(func)
@@ -107,6 +118,10 @@ def wrapper(*args, **kwds):
             return result
         return wrapper
 
+    if PY35:
+        import line_profiler_py35
+        wrap_coroutine = line_profiler_py35.wrap_coroutine
+
     def dump_stats(self, filename):
         """ Dump a representation of the data to a file as a pickled LineStats
         object from `get_stats()`.

line_profiler_py35.py

@@ -0,0 +1,16 @@
+""" This file is only imported in python 3.5 environments """
+import functools
+
+def wrap_coroutine(self, func):
+    """
+    Wrap a Python 3.5 coroutine to profile it.
+    """
+    @functools.wraps(func)
+    async def wrapper(*args, **kwds):
+        self.enable_by_count()
+        try:
+            result = await func(*args, **kwds)
+        finally:
+            self.disable_by_count()
+        return result
+    return wrapper

setup.py

@@ -1,4 +1,5 @@
 import os
+import sys
 
 # Monkeypatch distutils.
 import setuptools
@@ -33,6 +34,11 @@
 function-level profiling tools in the Python standard library.
 """
 
+
+py_modules = ['line_profiler', 'kernprof']
+if sys.version_info > (3, 4):
+    py_modules += ['line_profiler_py35']
+
 setup(
     name = 'line_profiler',
     version = '2.0',
@@ -66,7 +72,7 @@
         'Programming Language :: Python :: Implementation :: CPython',
         "Topic :: Software Development",
     ],
-    py_modules = ['line_profiler', 'kernprof'],
+    py_modules = py_modules,
     entry_points = {
         'console_scripts': [
             'kernprof=kernprof:main',

tests/_test_kernprof_py35.py

@@ -0,0 +1,23 @@
+from kernprof import ContextualProfile
+
+def test_coroutine_decorator(self):
+    async def _():
+        async def c(x):
+            """ A coroutine. """
+            y = x + 10
+            return y
+
+        profile = ContextualProfile()
+        c_wrapped = profile(c)
+        self.assertEqual(c_wrapped.__name__, c.__name__)
+        self.assertEqual(c_wrapped.__doc__, c.__doc__)
+
+        self.assertEqual(profile.enable_count, 0)
+        value = await c_wrapped(10)
+        self.assertEqual(profile.enable_count, 0)
+        self.assertEqual(value, await c(10))
+
+    import asyncio
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(_())
+    loop.close()

tests/test_kernprof.py

@@ -1,7 +1,11 @@
 import unittest
+import sys
 
 from kernprof import ContextualProfile
 
+PY3 = sys.version_info[0] == 3
+PY35 = PY3 and sys.version_info[1] >= 5
+
 
 def f(x):
     """ A function. """
@@ -72,3 +76,10 @@ def test_gen_decorator(self):
         with self.assertRaises(StopIteration):
             next(i)
         self.assertEqual(profile.enable_count, 0)
+
+    if PY35:
+        import _test_kernprof_py35
+        test_coroutine_decorator = _test_kernprof_py35.test_coroutine_decorator
+
+if __name__ == '__main__':
+    unittest.main()