support return in generators for Py3-only users.

Author: radix

effect/_test_do_py3.py

@@ -0,0 +1,11 @@
+# This code only works in Python 3, so it's left out of test_do.py, to be
+# optionally imported.
+
+from effect import Constant, Effect
+from effect.do import do
+
+
+@do
+def py3_generator_with_return():
+    yield Effect(Constant(1))
+    return 2  # noqa

effect/do.py

@@ -32,6 +32,8 @@ def foo():
     yielded Effect will be passed back into the generator as the result of the
     ``yield`` expression. Yielded :func:`do_return` values will provide the
     ultimate result of the Effect that is returned by the decorated function.
+    Note that :func:`do_return` is only necessary for Python 2 compatibility;
+    return statements can be used directly in Python 3-only code.
 
     It's important to note that any generator function decorated by ``@do``
     will no longer return a generator, but instead it will return an Effect,
@@ -78,6 +80,9 @@ def do_return(val):
         @do
         def foo():
             yield do_return('hello')
+
+    If you're writing Python 3-only code, you don't need to use this function,
+    and can just use the `return` statement as normal.
     """
     return _ReturnSentinel(val)
 
@@ -88,7 +93,7 @@ def _do(result, generator, is_error):
             val = generator.throw(*result)
         else:
             val = generator.send(result)
-    except StopIteration:
+    except StopIteration as stop:
         # If the generator we're spinning directly raises StopIteration, we'll
         # treat it like returning None from the function. But there may be a
         # case where some other code is raising StopIteration up through this
@@ -98,7 +103,12 @@ def _do(result, generator, is_error):
         if tb.tb_next:
             raise
         else:
-            return None
+            # Python 3 allows you to use `return val` in a generator, which
+            # will be translated to a `StopIteration` with a `value` attribute
+            # set to the return value. So we'll return that value as the
+            # ultimate result of the effect. Python 2 doesn't have the 'value'
+            # attribute of StopIteration, so we'll fall back to None.
+            return getattr(stop, 'value', None)
     if type(val) is _ReturnSentinel:
         return val.result
     elif type(val) is Effect:

effect/test_do.py

@@ -2,6 +2,9 @@
 from functools import partial
 
 from py.test import raises as raises
+from py.test import mark
+
+import six
 
 from testtools import TestCase
 from testtools.matchers import raises as match_raises, MatchesException
@@ -150,3 +153,10 @@ def f():
     eff = f()
     with raises(StopIteration):
         perf(eff)
+
+
+@mark.skipif(not six.PY3, reason="Testing a Py3-specific feature")
+def test_py3_return():
+    from effect._test_do_py3 import py3_generator_with_return
+    eff = py3_generator_with_return()
+    assert perf(eff) == 2