Merge pull request #57 from smarie/master

Improved `decorate` so as to support generator callers.

1 files changed, 29 insertions, 3 deletions

diff --git a/src/decorator.py b/src/decorator.py
index fd0e9ef..4433abf 100644
--- a/src/decorator.py
+++ b/src/decorator.py
@@ -216,9 +216,19 @@ 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.
+    If the caller is a generator function, the resulting function
+    will be a generator function.
     """
     evaldict = dict(_call_=caller, _func_=func)
     es = ''
@@ -226,9 +236,25 @@ def decorate(func, caller, extras=()):
         ex = '_e%d_' % i
         evaldict[ex] = extra
         es += ex + ', '
-    fun = FunctionMaker.create(
-        func, "return _call_(_func_, %s%%(shortsignature)s)" % es,
-        evaldict, __wrapped__=func)
+
+    if not ('3.5' <= sys.version < '3.6'):
+        create_generator = isgeneratorfunction(caller)
+    else:
+        # With Python 3.5: apparently 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:
+        fun = FunctionMaker.create(
+            func, "for res in _call_(_func_, %s%%(shortsignature)s):\n"
+                  "    yield res" % es, evaldict, __wrapped__=func)
+    else:
+        fun = FunctionMaker.create(
+            func, "return _call_(_func_, %s%%(shortsignature)s)" % es,
+            evaldict, __wrapped__=func)
     if hasattr(func, '__qualname__'):
         fun.__qualname__ = func.__qualname__
     return fun