only interpret StopIteration directly from the wrapped generator, not from things that the generator calls

Author: radix

effect/do.py

@@ -7,6 +7,7 @@
 
 from __future__ import print_function
 
+import sys
 import types
 
 from . import Effect, Func
@@ -88,7 +89,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)