gh-119824: Print stack entry when user input is needed (#119882)

Co-authored-by: Irit Katriel <1055913+iritkatriel@users.noreply.github.com>

1 files changed, 39 insertions, 11 deletions

diff --git a/Lib/pdb.py b/Lib/pdb.py
index ba84a29..ddbfb9d 100644
--- a/Lib/pdb.py
+++ b/Lib/pdb.py
@@ -603,10 +603,18 @@ class Pdb(bdb.Bdb, cmd.Cmd):
             assert tb is not None, "main exception must have a traceback"
         with self._hold_exceptions(_chained_exceptions):
             self.setup(frame, tb)
-            # if we have more commands to process, do not show the stack entry
-            if not self.cmdqueue:
+            # We should print the stack entry if and only if the user input
+            # is expected, and we should print it right before the user input.
+            # If self.cmdqueue is not empty, we append a "w 0" command to the
+            # queue, which is equivalent to print_stack_entry
+            if self.cmdqueue:
+                self.cmdqueue.append('w 0')
+            else:
                 self.print_stack_entry(self.stack[self.curindex])
             self._cmdloop()
+            # If "w 0" is not used, pop it out
+            if self.cmdqueue and self.cmdqueue[-1] == 'w 0':
+                self.cmdqueue.pop()
             self.forget()
 
     def displayhook(self, obj):
@@ -1401,16 +1409,24 @@ class Pdb(bdb.Bdb, cmd.Cmd):
     complete_cl = _complete_location
 
     def do_where(self, arg):
-        """w(here)
+        """w(here) [count]
 
-        Print a stack trace, with the most recent frame at the bottom.
+        Print a stack trace. If count is not specified, print the full stack.
+        If count is 0, print the current frame entry. If count is positive,
+        print count entries from the most recent frame. If count is negative,
+        print -count entries from the least recent frame.
         An arrow indicates the "current frame", which determines the
         context of most commands.  'bt' is an alias for this command.
         """
-        if arg:
-            self._print_invalid_arg(arg)
-            return
-        self.print_stack_trace()
+        if not arg:
+            count = None
+        else:
+            try:
+                count = int(arg)
+            except ValueError:
+                self.error('Invalid count (%s)' % arg)
+                return
+        self.print_stack_trace(count)
     do_w = do_where
     do_bt = do_where
 
@@ -2065,10 +2081,22 @@ class Pdb(bdb.Bdb, cmd.Cmd):
     # It is also consistent with the up/down commands (which are
     # compatible with dbx and gdb: up moves towards 'main()'
     # and down moves towards the most recent stack frame).
-
-    def print_stack_trace(self):
+    #     * if count is None, prints the full stack
+    #     * if count = 0, prints the current frame entry
+    #     * if count < 0, prints -count least recent frame entries
+    #     * if count > 0, prints count most recent frame entries
+
+    def print_stack_trace(self, count=None):
+        if count is None:
+            stack_to_print = self.stack
+        elif count == 0:
+            stack_to_print = [self.stack[self.curindex]]
+        elif count < 0:
+            stack_to_print = self.stack[:-count]
+        else:
+            stack_to_print = self.stack[-count:]
         try:
-            for frame_lineno in self.stack:
+            for frame_lineno in stack_to_print:
                 self.print_stack_entry(frame_lineno)
         except KeyboardInterrupt:
             pass