diff --git a/line_profiler.py b/line_profiler.py index aac01c8..960e67d 100755 --- a/line_profiler.py +++ b/line_profiler.py @@ -22,6 +22,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: @@ -41,6 +42,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 @@ -60,7 +69,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) @@ -102,6 +113,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()`. diff --git a/line_profiler_py35.py b/line_profiler_py35.py new file mode 100644 index 0000000..c400094 --- /dev/null +++ b/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 diff --git a/setup.py b/setup.py index a8042c0..c0272de 100755 --- a/setup.py +++ b/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 = '1.0', @@ -64,7 +70,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', diff --git a/tests/_test_kernprof_py35.py b/tests/_test_kernprof_py35.py new file mode 100644 index 0000000..c52c8cf --- /dev/null +++ b/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() diff --git a/tests/test_kernprof.py b/tests/test_kernprof.py index d230428..7702b54 100644 --- a/tests/test_kernprof.py +++ b/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()