Merge pull request #56 from python-effect/local-only-stopiteration

re-raise StopIteration unless it comes directly from the @do-wrapped generator

Author: radix

effect/do.py

@@ -7,6 +7,7 @@
 
 from __future__ import print_function
 
+import sys
 import types
 
 from . import Effect, Func
@@ -26,11 +27,12 @@ def foo():
         eff = foo()
         return eff.on(...)
 
-    ``@do`` must decorate a generator function. Any yielded values must either
-    be Effects or the result of a :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.
+    ``@do`` must decorate a generator function (not any other type of
+    iterator). Any yielded values must either be Effects or the result of a
+    :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.
 
     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,
@@ -88,7 +90,16 @@ def _do(result, generator, is_error):
         else:
             val = generator.send(result)
     except StopIteration:
-        return None
+        # 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
+        # generator, in which case we shouldn't really treat it like a function
+        # return -- it could quite easily hide bugs.
+        tb = sys.exc_info()[2]
+        if tb.tb_next:
+            raise
+        else:
+            return None
     if type(val) is _ReturnSentinel:
         return val.result
     elif type(val) is Effect:

effect/test_do.py

@@ -1,8 +1,10 @@
 import sys
 from functools import partial
 
+from py.test import raises as raises
+
 from testtools import TestCase
-from testtools.matchers import raises, MatchesException
+from testtools.matchers import raises as match_raises, MatchesException
 
 from . import (
     ComposedDispatcher, Constant, Effect, Error, TypeDispatcher,
@@ -21,7 +23,7 @@ def test_do_non_gf(self):
         f = lambda: None
         self.assertThat(
             lambda: perf(do(f)()),
-            raises(TypeError(
+            match_raises(TypeError(
                 "%r is not a generator function. It returned None." % (f,)
             )))
 
@@ -133,3 +135,18 @@ def f():
         eff = f()
         self.assertEqual(perf(eff), 'foo')
         self.assertEqual(perf(eff), 'foo')
+
+
+def test_stop_iteration_only_local():
+    """
+    Arbitrary :obj:`StopIteration` exceptions are not treated the same way as
+    falling off the end of the generator -- they are raised through.
+    """
+    @do
+    def f():
+        raise StopIteration()
+        yield Effect(Constant('foo'))
+
+    eff = f()
+    with raises(StopIteration):
+        perf(eff)