Add support for generators

pysnooper/pysnooper.py

@@ -33,8 +33,7 @@ def write(s):
     return write
 
 
-
-def snoop(output=None, variables=(), depth=1, prefix=''):
+def snoop(output=None, variables=(), depth=1, prefix='', is_generator=False):
     '''
     Snoop on the function, writing everything it's doing to stderr.
 
@@ -69,8 +68,20 @@ def decorate(function, *args, **kwargs):
         with Tracer(target_code_object=target_code_object,
                     write=write, variables=variables,
                     depth=depth, prefix=prefix):
-            return function(*args, **kwargs)
+                return function(*args, **kwargs)
+
+    @decorator.decorator
+    def generator_decorate(function, *args, **kwargs):
+        target_code_object = function.__code__
+        with Tracer(target_code_object=target_code_object,
+                    write=write, variables=variables,
+                    depth=depth, prefix=prefix):
+                for v in function(*args, **kwargs):
+                    yield v
+                # yield from function(*args, **kwargs)
 
+    if is_generator:
+        return generator_decorate
     return decorate
 
 

tests/test_pysnooper.py

@@ -1,6 +1,7 @@
 # Copyright 2019 Ram Rachum.
 # This program is distributed under the MIT license.
 
+import inspect
 import io
 import re
 import abc
@@ -197,3 +198,33 @@ def my_function(foo):
             )
         )
 
+
+def test_generator():
+    string_io = io.StringIO()
+
+    @pysnooper.snoop(string_io)
+    def generator_empty_output():
+        yield 1
+
+    result = generator_empty_output()
+    assert inspect.isgenerator(result)
+    assert next(result) == 1
+    output = string_io.getvalue()
+    assert_output(output, ())
+
+    @pysnooper.snoop(string_io, is_generator=True)
+    def generator():
+        yield 1
+
+    result = generator()
+    assert inspect.isgenerator(result)
+    assert next(result) == 1
+    output = string_io.getvalue()
+    assert_output(
+        output,
+        (
+            CallEntry(),
+            LineEntry('yield 1'),
+            ReturnEntry('yield 1'),
+        )
+    )