Added a test for decorating generators

3 files changed, 36 insertions, 17 deletions

diff --git a/CHANGES.md b/CHANGES.md
index 00db976..83b49af 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,6 +1,12 @@
 HISTORY
 --------
 
+## unreleased
+
+Accepted a patch from Sylvain Marie (https://github.com/smarie): now the
+decorator module can decorate generator functions by preserving their
+being generator functions.
+
 ## 4.3.1 (2018-08-04)
 
 Added a section "For the impatient" to the README, addressing an issue
diff --git a/src/decorator.py b/src/decorator.py
index 4433abf..80507ed 100644
--- a/src/decorator.py
+++ b/src/decorator.py
@@ -65,6 +65,12 @@ except AttributeError:
     # let's assume there are no coroutine functions in old Python
     def iscoroutinefunction(f):
         return False
+try:
+    from inspect import isgeneratorfunction
+except ImportError:
+    # assume no generator function in old Python versions
+    def isgeneratorfunction():
+        return False
 
 
 DEF = re.compile(r'\s*def\s*([_\w][_\w\d]*)\s*\(')
@@ -216,14 +222,6 @@ class FunctionMaker(object):
         return self.make(body, evaldict, addsource, **attrs)
 
 
-try:
-    from inspect import isgeneratorfunction
-except ImportError:
-    # assume no generator function in very old python versions
-    def isgeneratorfunction():
-        return False
-
-
 def decorate(func, caller, extras=()):
     """
     decorate(func, caller) decorates a function using a caller.
@@ -237,17 +235,15 @@ def decorate(func, caller, extras=()):
         evaldict[ex] = extra
         es += ex + ', '
 
-    if not ('3.5' <= sys.version < '3.6'):
-        create_generator = isgeneratorfunction(caller)
-    else:
-        # With Python 3.5: apparently isgeneratorfunction returns
-        # True for all coroutines
-
+    if sys.version_info[2:] == (3, 5):
+        # With Python 3.5 isgeneratorfunction returns True for all coroutines
         # However we know that it is NOT possible to have a generator
         # coroutine in python 3.5: PEP525 was not there yet.
-        create_generator = isgeneratorfunction(caller) and not iscoroutinefunction(caller)
-
-    if create_generator:
+        generatorcaller = isgeneratorfunction(
+            caller) and not iscoroutinefunction(caller)
+    else:
+        generatorcaller = isgeneratorfunction(caller)
+    if generatorcaller:
         fun = FunctionMaker.create(
             func, "for res in _call_(_func_, %s%%(shortsignature)s):\n"
                   "    yield res" % es, evaldict, __wrapped__=func)
diff --git a/src/tests/test.py b/src/tests/test.py
index 3c4dcd1..9068b11 100644
--- a/src/tests/test.py
+++ b/src/tests/test.py
@@ -58,6 +58,23 @@ class CoroutineTestCase(unittest.TestCase):
 ''')
 
 
+def gen123():
+    yield 1
+    yield 2
+    yield 3
+
+
+class GeneratorCallerTestCase(unittest.TestCase):
+    def test_gen123(self):
+        @decorator
+        def square(func, *args, **kw):
+            for x in gen123():
+                yield x * x
+        new = square(gen123)
+        self.assertTrue(inspect.isgeneratorfunction(new))
+        self.assertEqual(list(new()), [1, 4, 9])
+
+
 class DocumentationTestCase(unittest.TestCase):
     def test(self):
         err = doctest.testmod(doc)[0]