Merge pull request #58 from python-effect/py3-stopiteration-result

Allow `return x` in @Do generators for Python 3 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,7 +32,9 @@ def foo():
     :func:`do_return` call. The result of a 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.
+    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,
@@ -79,6 +81,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)
 
@@ -89,7 +94,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
@@ -99,7 +104,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,11 @@ def f():
     eff = f()
     with raises(StopIteration):
         perf(eff)
+
+
+@mark.skipif(not six.PY3, reason="Testing a Py3-specific feature")
+def test_py3_return():
+    """The `return x` syntax in Py3 sets the result of the Effect to `x`."""
+    from effect._test_do_py3 import py3_generator_with_return
+    eff = py3_generator_with_return()
+    assert perf(eff) == 2