Better behavior when stepping over yield[from]. Fixes issue 16596. By Xavier de Gaye.

4 files changed, 351 insertions, 7 deletions

diff --git a/Lib/bdb.py b/Lib/bdb.py
index dd1f4287..67a0846 100644
--- a/Lib/bdb.py
+++ b/Lib/bdb.py
@@ -3,6 +3,7 @@
 import fnmatch
 import sys
 import os
+from inspect import CO_GENERATOR
 
 __all__ = ["BdbQuit", "Bdb", "Breakpoint"]
 
@@ -75,24 +76,48 @@ class Bdb:
         if not (self.stop_here(frame) or self.break_anywhere(frame)):
             # No need to trace this function
             return # None
+        # Ignore call events in generator except when stepping.
+        if self.stopframe and frame.f_code.co_flags & CO_GENERATOR:
+            return self.trace_dispatch
         self.user_call(frame, arg)
         if self.quitting: raise BdbQuit
         return self.trace_dispatch
 
     def dispatch_return(self, frame, arg):
         if self.stop_here(frame) or frame == self.returnframe:
+            # Ignore return events in generator except when stepping.
+            if self.stopframe and frame.f_code.co_flags & CO_GENERATOR:
+                return self.trace_dispatch
             try:
                 self.frame_returning = frame
                 self.user_return(frame, arg)
             finally:
                 self.frame_returning = None
             if self.quitting: raise BdbQuit
+            # The user issued a 'next' or 'until' command.
+            if self.stopframe is frame and self.stoplineno != -1:
+                self._set_stopinfo(None, None)
         return self.trace_dispatch
 
     def dispatch_exception(self, frame, arg):
         if self.stop_here(frame):
+            # When stepping with next/until/return in a generator frame, skip
+            # the internal StopIteration exception (with no traceback)
+            # triggered by a subiterator run with the 'yield from' statement.
+            if not (frame.f_code.co_flags & CO_GENERATOR
+                    and arg[0] is StopIteration and arg[2] is None):
+                self.user_exception(frame, arg)
+                if self.quitting: raise BdbQuit
+        # Stop at the StopIteration or GeneratorExit exception when the user
+        # has set stopframe in a generator by issuing a return command, or a
+        # next/until command at the last statement in the generator before the
+        # exception.
+        elif (self.stopframe and frame is not self.stopframe
+                and self.stopframe.f_code.co_flags & CO_GENERATOR
+                and arg[0] in (StopIteration, GeneratorExit)):
             self.user_exception(frame, arg)
             if self.quitting: raise BdbQuit
+
         return self.trace_dispatch
 
     # Normally derived classes don't override the following
@@ -115,10 +140,8 @@ class Bdb:
             if self.stoplineno == -1:
                 return False
             return frame.f_lineno >= self.stoplineno
-        while frame is not None and frame is not self.stopframe:
-            if frame is self.botframe:
-                return True
-            frame = frame.f_back
+        if not self.stopframe:
+            return True
         return False
 
     def break_here(self, frame):
@@ -207,7 +230,10 @@ class Bdb:
 
     def set_return(self, frame):
         """Stop when returning from the given frame."""
-        self._set_stopinfo(frame.f_back, frame)
+        if frame.f_code.co_flags & CO_GENERATOR:
+            self._set_stopinfo(frame, None, -1)
+        else:
+            self._set_stopinfo(frame.f_back, frame)
 
     def set_trace(self, frame=None):
         """Start debugging from `frame`.
diff --git a/Lib/pdb.py b/Lib/pdb.py
index 2268d30..dd7ceb8 100755
--- a/Lib/pdb.py
+++ b/Lib/pdb.py
@@ -297,8 +297,16 @@ class Pdb(bdb.Bdb, cmd.Cmd):
             return
         exc_type, exc_value, exc_traceback = exc_info
         frame.f_locals['__exception__'] = exc_type, exc_value
-        self.message(traceback.format_exception_only(exc_type,
-                                                     exc_value)[-1].strip())
+
+        # An 'Internal StopIteration' exception is an exception debug event
+        # issued by the interpreter when handling a subgenerator run with
+        # 'yield from' or a generator controled by a for loop. No exception has
+        # actually occured in this case. The debugger uses this debug event to
+        # stop when the debuggee is returning from such generators.
+        prefix = 'Internal ' if (not exc_traceback
+                                    and exc_type is StopIteration) else ''
+        self.message('%s%s' % (prefix,
+            traceback.format_exception_only(exc_type, exc_value)[-1].strip()))
         self.interaction(frame, exc_traceback)
 
     # General interaction function
diff --git a/Lib/test/test_pdb.py b/Lib/test/test_pdb.py
index 7993d02..74253b3 100644
--- a/Lib/test/test_pdb.py
+++ b/Lib/test/test_pdb.py
@@ -600,6 +600,311 @@ def test_pdb_run_with_code_object():
     (Pdb) continue
     """
 
+def test_next_until_return_at_return_event():
+    """Test that pdb stops after a next/until/return issued at a return debug event.
+
+    >>> def test_function_2():
+    ...     x = 1
+    ...     x = 2
+
+    >>> def test_function():
+    ...     import pdb; pdb.Pdb(nosigint=True).set_trace()
+    ...     test_function_2()
+    ...     test_function_2()
+    ...     test_function_2()
+    ...     end = 1
+
+    >>> with PdbTestInput(['break test_function_2',
+    ...                    'continue',
+    ...                    'return',
+    ...                    'next',
+    ...                    'continue',
+    ...                    'return',
+    ...                    'until',
+    ...                    'continue',
+    ...                    'return',
+    ...                    'return',
+    ...                    'continue']):
+    ...     test_function()
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[1]>(3)test_function()
+    -> test_function_2()
+    (Pdb) break test_function_2
+    Breakpoint 1 at <doctest test.test_pdb.test_next_until_return_at_return_event[0]>:1
+    (Pdb) continue
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(2)test_function_2()
+    -> x = 1
+    (Pdb) return
+    --Return--
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(3)test_function_2()->None
+    -> x = 2
+    (Pdb) next
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[1]>(4)test_function()
+    -> test_function_2()
+    (Pdb) continue
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(2)test_function_2()
+    -> x = 1
+    (Pdb) return
+    --Return--
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(3)test_function_2()->None
+    -> x = 2
+    (Pdb) until
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[1]>(5)test_function()
+    -> test_function_2()
+    (Pdb) continue
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(2)test_function_2()
+    -> x = 1
+    (Pdb) return
+    --Return--
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[0]>(3)test_function_2()->None
+    -> x = 2
+    (Pdb) return
+    > <doctest test.test_pdb.test_next_until_return_at_return_event[1]>(6)test_function()
+    -> end = 1
+    (Pdb) continue
+    """
+
+def test_pdb_next_command_for_generator():
+    """Testing skip unwindng stack on yield for generators for "next" command
+
+    >>> def test_gen():
+    ...     yield 0
+    ...     return 1
+    ...     yield 2
+
+    >>> def test_function():
+    ...     import pdb; pdb.Pdb(nosigint=True).set_trace()
+    ...     it = test_gen()
+    ...     try:
+    ...         assert next(it) == 0
+    ...         next(it)
+    ...     except StopIteration as ex:
+    ...         assert ex.value == 1
+    ...     print("finished")
+
+    >>> with PdbTestInput(['step',
+    ...                    'step',
+    ...                    'step',
+    ...                    'next',
+    ...                    'next',
+    ...                    'step',
+    ...                    'step',
+    ...                    'continue']):
+    ...     test_function()
+    > <doctest test.test_pdb.test_pdb_next_command_for_generator[1]>(3)test_function()
+    -> it = test_gen()
+    (Pdb) step
+    > <doctest test.test_pdb.test_pdb_next_command_for_generator[1]>(4)test_function()
+    -> try:
+    (Pdb) step
+    > <doctest test.test_pdb.test_pdb_next_command_for_generator[1]>(5)test_function()
+    -> assert next(it) == 0
+    (Pdb) step
+    --Call--
+    > <doctest test.test_pdb.test_pdb_next_command_for_generator[0]>(1)test_gen()
+    -> def test_gen():
+    (Pdb) next
+    > <doctest test.test_pdb.test_pdb_next_command_for_generator[0]>(2)test_gen()
+    -> yield 0
+    (Pdb) next
+    > <doctest test.test_pdb.test_pdb_next_command_for_generator[0]>(3)test_gen()
+    -> return 1
+    (Pdb) step
+    --Return--
+    > <doctest test.test_pdb.test_pdb_next_command_for_generator[0]>(3)test_gen()->1
+    -> return 1
+    (Pdb) step
+    StopIteration: 1
+    > <doctest test.test_pdb.test_pdb_next_command_for_generator[1]>(6)test_function()
+    -> next(it)
+    (Pdb) continue
+    finished
+    """
+
+def test_pdb_return_command_for_generator():
+    """Testing no unwindng stack on yield for generators
+       for "return" command
+
+    >>> def test_gen():
+    ...     yield 0
+    ...     return 1
+    ...     yield 2
+
+    >>> def test_function():
+    ...     import pdb; pdb.Pdb(nosigint=True).set_trace()
+    ...     it = test_gen()
+    ...     try:
+    ...         assert next(it) == 0
+    ...         next(it)
+    ...     except StopIteration as ex:
+    ...         assert ex.value == 1
+    ...     print("finished")
+
+    >>> with PdbTestInput(['step',
+    ...                    'step',
+    ...                    'step',
+    ...                    'return',
+    ...                    'step',
+    ...                    'step',
+    ...                    'continue']):
+    ...     test_function()
+    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(3)test_function()
+    -> it = test_gen()
+    (Pdb) step
+    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(4)test_function()
+    -> try:
+    (Pdb) step
+    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(5)test_function()
+    -> assert next(it) == 0
+    (Pdb) step
+    --Call--
+    > <doctest test.test_pdb.test_pdb_return_command_for_generator[0]>(1)test_gen()
+    -> def test_gen():
+    (Pdb) return
+    StopIteration: 1
+    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(6)test_function()
+    -> next(it)
+    (Pdb) step
+    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(7)test_function()
+    -> except StopIteration as ex:
+    (Pdb) step
+    > <doctest test.test_pdb.test_pdb_return_command_for_generator[1]>(8)test_function()
+    -> assert ex.value == 1
+    (Pdb) continue
+    finished
+    """
+
+def test_pdb_until_command_for_generator():
+    """Testing no unwindng stack on yield for generators
+       for "until" command if target breakpoing is not reached
+
+    >>> def test_gen():
+    ...     yield 0
+    ...     yield 1
+    ...     yield 2
+
+    >>> def test_function():
+    ...     import pdb; pdb.Pdb(nosigint=True).set_trace()
+    ...     for i in test_gen():
+    ...         print(i)
+    ...     print("finished")
+
+    >>> with PdbTestInput(['step',
+    ...                    'until 4',
+    ...                    'step',
+    ...                    'step',
+    ...                    'continue']):
+    ...     test_function()
+    > <doctest test.test_pdb.test_pdb_until_command_for_generator[1]>(3)test_function()
+    -> for i in test_gen():
+    (Pdb) step
+    --Call--
+    > <doctest test.test_pdb.test_pdb_until_command_for_generator[0]>(1)test_gen()
+    -> def test_gen():
+    (Pdb) until 4
+    0
+    1
+    > <doctest test.test_pdb.test_pdb_until_command_for_generator[0]>(4)test_gen()
+    -> yield 2
+    (Pdb) step
+    --Return--
+    > <doctest test.test_pdb.test_pdb_until_command_for_generator[0]>(4)test_gen()->2
+    -> yield 2
+    (Pdb) step
+    > <doctest test.test_pdb.test_pdb_until_command_for_generator[1]>(4)test_function()
+    -> print(i)
+    (Pdb) continue
+    2
+    finished
+    """
+
+def test_pdb_next_command_in_generator_for_loop():
+    """The next command on returning from a generator controled by a for loop.
+
+    >>> def test_gen():
+    ...     yield 0
+    ...     return 1
+
+    >>> def test_function():
+    ...     import pdb; pdb.Pdb(nosigint=True).set_trace()
+    ...     for i in test_gen():
+    ...         print('value', i)
+    ...     x = 123
+
+    >>> with PdbTestInput(['break test_gen',
+    ...                    'continue',
+    ...                    'next',
+    ...                    'next',
+    ...                    'next',
+    ...                    'continue']):
+    ...     test_function()
+    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[1]>(3)test_function()
+    -> for i in test_gen():
+    (Pdb) break test_gen
+    Breakpoint 6 at <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[0]>:1
+    (Pdb) continue
+    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[0]>(2)test_gen()
+    -> yield 0
+    (Pdb) next
+    value 0
+    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[0]>(3)test_gen()
+    -> return 1
+    (Pdb) next
+    Internal StopIteration: 1
+    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[1]>(3)test_function()
+    -> for i in test_gen():
+    (Pdb) next
+    > <doctest test.test_pdb.test_pdb_next_command_in_generator_for_loop[1]>(5)test_function()
+    -> x = 123
+    (Pdb) continue
+    """
+
+def test_pdb_next_command_subiterator():
+    """The next command in a generator with a subiterator.
+
+    >>> def test_subgenerator():
+    ...     yield 0
+    ...     return 1
+
+    >>> def test_gen():
+    ...     x = yield from test_subgenerator()
+    ...     return x
+
+    >>> def test_function():
+    ...     import pdb; pdb.Pdb(nosigint=True).set_trace()
+    ...     for i in test_gen():
+    ...         print('value', i)
+    ...     x = 123
+
+    >>> with PdbTestInput(['step',
+    ...                    'step',
+    ...                    'next',
+    ...                    'next',
+    ...                    'next',
+    ...                    'continue']):
+    ...     test_function()
+    > <doctest test.test_pdb.test_pdb_next_command_subiterator[2]>(3)test_function()
+    -> for i in test_gen():
+    (Pdb) step
+    --Call--
+    > <doctest test.test_pdb.test_pdb_next_command_subiterator[1]>(1)test_gen()
+    -> def test_gen():
+    (Pdb) step
+    > <doctest test.test_pdb.test_pdb_next_command_subiterator[1]>(2)test_gen()
+    -> x = yield from test_subgenerator()
+    (Pdb) next
+    value 0
+    > <doctest test.test_pdb.test_pdb_next_command_subiterator[1]>(3)test_gen()
+    -> return x
+    (Pdb) next
+    Internal StopIteration: 1
+    > <doctest test.test_pdb.test_pdb_next_command_subiterator[2]>(3)test_function()
+    -> for i in test_gen():
+    (Pdb) next
+    > <doctest test.test_pdb.test_pdb_next_command_subiterator[2]>(5)test_function()
+    -> x = 123
+    (Pdb) continue
+    """
+
 
 class PdbTestCase(unittest.TestCase):
 
diff --git a/Python/ceval.c b/Python/ceval.c
index e2e882b..ab419dc 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -1904,6 +1904,9 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             Py_DECREF(v);
             if (retval == NULL) {
                 PyObject *val;
+                if (tstate->c_tracefunc != NULL
+                        && PyErr_ExceptionMatches(PyExc_StopIteration))
+                    call_exc_trace(tstate->c_tracefunc, tstate->c_traceobj, f);
                 err = _PyGen_FetchStopIterationValue(&val);
                 if (err < 0)
                     goto error;
@@ -2654,6 +2657,8 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             if (PyErr_Occurred()) {
                 if (!PyErr_ExceptionMatches(PyExc_StopIteration))
                     goto error;
+                else if (tstate->c_tracefunc != NULL)
+                    call_exc_trace(tstate->c_tracefunc, tstate->c_traceobj, f);
                 PyErr_Clear();
             }
             /* iterator ended normally */