Add support for wrapping generators/coroutines/asyncgenerators with hunter.wrap.

Author: ionelmc

pytest.ini

@@ -22,8 +22,6 @@ addopts =
     --ignore=setup.py
     --ignore=ci
     --ignore=.eggs
-    --doctest-modules
-    --doctest-glob=\*.rst
     --tb=short
 testpaths =
     tests

src/hunter/__init__.py

@@ -469,29 +469,63 @@ def my_function():
     """
 
     def tracing_decorator(func):
-        @functools.wraps(func)
-        def tracing_wrapper(*args, **kwargs):
-            predicates = []
-            local = trace_options.pop('local', False)
-            if local:
-                predicates.append(Query(depth_lt=2))
-            predicates.append(
-                From(
-                    Query(kind='call'),
-                    Not(
-                        When(
-                            Query(calls_gt=0, depth=0) & Not(Query(kind='return')),
-                            Stop,
-                        )
-                    ),
-                    watermark=-1,
-                )
+        predicates = []
+        local = trace_options.pop('local', False)
+        if local:
+            predicates.append(Query(depth_lt=2))
+        predicates.append(
+            From(
+                Query(kind='call', function=func.__name__),
+                Not(
+                    When(
+                        Query(calls_gt=0, depth=0) & Not(Query(kind='return')),
+                        Stop,
+                    )
+                ),
+                watermark=-1,
             )
-            local_tracer = trace(*predicates, **trace_options)
-            try:
-                return func(*args, **kwargs)
-            finally:
-                local_tracer.stop()
+        )
+
+        if inspect.isasyncgenfunction(func):
+
+            @functools.wraps(func)
+            async def tracing_wrapper(*args, **kwargs):
+                local_tracer = trace(*predicates, **trace_options)
+                try:
+                    async for item in func(*args, **kwargs):
+                        yield item
+                finally:
+                    local_tracer.stop()
+
+        elif inspect.iscoroutinefunction(func):
+
+            @functools.wraps(func)
+            async def tracing_wrapper(*args, **kwargs):
+                local_tracer = trace(*predicates, **trace_options)
+                try:
+                    return await func(*args, **kwargs)
+                finally:
+                    local_tracer.stop()
+
+        elif inspect.isgeneratorfunction(func):
+
+            @functools.wraps(func)
+            def tracing_wrapper(*args, **kwargs):
+                local_tracer = trace(*predicates, **trace_options)
+                try:
+                    yield from func(*args, **kwargs)
+                finally:
+                    local_tracer.stop()
+
+        else:
+
+            @functools.wraps(func)
+            def tracing_wrapper(*args, **kwargs):
+                local_tracer = trace(*predicates, **trace_options)
+                try:
+                    return func(*args, **kwargs)
+                finally:
+                    local_tracer.stop()
 
         return tracing_wrapper
 

tests/test_tracer.py

@@ -1,5 +1,6 @@
 from __future__ import print_function
 
+import asyncio
 import functools
 import os
 import platform
@@ -35,7 +36,6 @@
 except ImportError:
     from io import StringIO
 
-
 if hunter.Tracer.__module__ == 'hunter.tracer':
 
     class EvilFrame(object):
@@ -85,7 +85,6 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 else:
     from eviltracer import EvilTracer
 
-
 trace = EvilTracer
 
 pytest_plugins = ('pytester',)
@@ -599,6 +598,89 @@ def foo():
         assert 'tracer.stop()' not in call
 
 
+def test_wraps_generator(LineMatcher):
+    calls = []
+
+    @hunter.wrap(action=lambda event: calls.append('%6r calls=%r depth=%r %s' % (event.kind, event.calls, event.depth, event.fullsource)))
+    def foo():
+        yield 1
+
+    assert list(foo()) == [1]
+    lm = LineMatcher(calls)
+    for line in calls:
+        print(repr(line))
+    lm.fnmatch_lines(
+        [
+            "'call' calls=0 depth=0     @hunter.wrap*",
+            "'line' calls=1 depth=1         yield 1\n",
+            "'return' calls=1 depth=0         yield 1\n",
+        ]
+    )
+    for call in calls:
+        assert 'tracer.stop()' not in call
+
+
+def test_wraps_async_generator(LineMatcher):
+    calls = []
+
+    @hunter.wrap(
+        action=lambda event: calls.append(
+            '%s %6r calls=%r depth=%r %s' % (event.module, event.kind, event.calls, event.depth, event.fullsource)
+        )
+    )
+    async def foo():
+        yield 1
+        await asyncio.sleep(0.01)
+
+    async def runner():
+        result = []
+        async for item in foo():
+            result.append(item)
+        return result
+
+    assert asyncio.run(runner()) == [1]
+    lm = LineMatcher(calls)
+    for line in calls:
+        print(repr(line))
+    lm.fnmatch_lines(
+        [
+            "test_tracer 'call' calls=* depth=0     @hunter.wrap*",
+            "test_tracer 'line' calls=* depth=1         yield 1\n",
+            "test_tracer 'return' calls=* depth=0         yield 1\n",
+        ]
+    )
+    for call in calls:
+        assert 'tracer.stop()' not in call
+
+
+def test_wraps_coroutine(LineMatcher):
+    calls = []
+
+    @hunter.wrap(
+        action=lambda event: calls.append(
+            '%s %6r calls=%r depth=%r %s' % (event.module, event.kind, event.calls, event.depth, event.fullsource)
+        )
+    )
+    async def foo():
+        await asyncio.sleep(0.01)
+        print(1)
+        return 1
+
+    assert asyncio.run(foo()) == 1
+    lm = LineMatcher(calls)
+    for line in calls:
+        print(repr(line))
+    lm.fnmatch_lines(
+        [
+            "test_tracer 'call' calls=0 depth=0     @hunter.wrap*",
+            "test_tracer 'line' calls=1 depth=1         await asyncio.sleep(0.01)\n",
+            "test_tracer 'return' calls=* depth=0         await asyncio.sleep(0.01)\n",
+        ]
+    )
+    for call in calls:
+        assert 'tracer.stop()' not in call
+
+
 def test_wraps_local(LineMatcher):
     calls = []
 
@@ -630,6 +712,111 @@ def foo():
         assert 'tracer.stop()' not in call
 
 
+def test_wraps_generator_local(LineMatcher):
+    calls = []
+
+    def bar():
+        for i in range(2):
+            return 'A'
+
+    @hunter.wrap(
+        local=True, action=lambda event: calls.append('%6r calls=%r depth=%r %s' % (event.kind, event.calls, event.depth, event.fullsource))
+    )
+    def foo():
+        bar()
+        yield 1
+
+    assert list(foo()) == [1]
+    lm = LineMatcher(calls)
+    for line in calls:
+        print(repr(line))
+    lm.fnmatch_lines(
+        [
+            "'call' calls=0 depth=0     @hunter.wrap*",
+            "'line' calls=* depth=1         yield 1\n",
+            "'return' calls=* depth=0         yield 1\n",
+        ]
+    )
+    for call in calls:
+        assert 'for i in range(2)' not in call
+        assert 'tracer.stop()' not in call
+
+
+def test_wraps_async_generator_local(LineMatcher):
+    calls = []
+
+    def bar():
+        for i in range(2):
+            return 'A'
+
+    @hunter.wrap(
+        local=True,
+        action=lambda event: calls.append(
+            '%s %6r calls=%r depth=%r %s' % (event.module, event.kind, event.calls, event.depth, event.fullsource)
+        ),
+    )
+    async def foo():
+        bar()
+        yield 1
+        await asyncio.sleep(0.01)
+
+    async def runner():
+        result = []
+        async for item in foo():
+            result.append(item)
+        return result
+
+    assert asyncio.run(runner()) == [1]
+    lm = LineMatcher(calls)
+    for line in calls:
+        print(repr(line))
+    lm.fnmatch_lines(
+        [
+            "test_tracer 'call' calls=* depth=0     @hunter.wrap*",
+            "test_tracer 'line' calls=* depth=1         yield 1\n",
+            "test_tracer 'return' calls=* depth=0         yield 1\n",
+        ]
+    )
+    for call in calls:
+        assert 'for i in range(2)' not in call
+        assert 'tracer.stop()' not in call
+
+
+def test_wraps_coroutine_local(LineMatcher):
+    calls = []
+
+    def bar():
+        for i in range(2):
+            return 'A'
+
+    @hunter.wrap(
+        local=True,
+        action=lambda event: calls.append(
+            '%s %6r calls=%r depth=%r %s' % (event.module, event.kind, event.calls, event.depth, event.fullsource)
+        ),
+    )
+    async def foo():
+        bar()
+        await asyncio.sleep(0.01)
+        print(1)
+        return 1
+
+    assert asyncio.run(foo()) == 1
+    lm = LineMatcher(calls)
+    for line in calls:
+        print(repr(line))
+    lm.fnmatch_lines(
+        [
+            "test_tracer 'call' calls=0 depth=0     @hunter.wrap*",
+            "test_tracer 'line' calls=* depth=1         await asyncio.sleep(0.01)\n",
+            "test_tracer 'return' calls=* depth=0         await asyncio.sleep(0.01)\n",
+        ]
+    )
+    for call in calls:
+        assert 'for i in range(2)' not in call
+        assert 'tracer.stop()' not in call
+
+
 @pytest.mark.skipif('os.environ.get("SETUPPY_CFLAGS") == "-DCYTHON_TRACE=1"')
 def test_depth():
     calls = []