gh-121468: Support async breakpoint in pdb (#132576)

Lib/pdb.py

@@ -385,6 +385,9 @@ class Pdb(bdb.Bdb, cmd.Cmd):
         self.commands_bnum = None # The breakpoint number for which we are
                                   # defining a list
 
+        self.async_shim_frame = None
+        self.async_awaitable = None
+
         self._chained_exceptions = tuple()
         self._chained_exception_index = 0
 
@@ -400,6 +403,57 @@ class Pdb(bdb.Bdb, cmd.Cmd):
 
         super().set_trace(frame)
 
+    async def set_trace_async(self, frame=None, *, commands=None):
+        if self.async_awaitable is not None:
+            # We are already in a set_trace_async call, do not mess with it
+            return
+
+        if frame is None:
+            frame = sys._getframe().f_back
+
+        # We need set_trace to set up the basics, however, this will call
+        # set_stepinstr() will we need to compensate for, because we don't
+        # want to trigger on calls
+        self.set_trace(frame, commands=commands)
+        # Changing the stopframe will disable trace dispatch on calls
+        self.stopframe = frame
+        # We need to stop tracing because we don't have the privilege to avoid
+        # triggering tracing functions as normal, as we are not already in
+        # tracing functions
+        self.stop_trace()
+
+        self.async_shim_frame = sys._getframe()
+        self.async_awaitable = None
+
+        while True:
+            self.async_awaitable = None
+            # Simulate a trace event
+            # This should bring up pdb and make pdb believe it's debugging the
+            # caller frame
+            self.trace_dispatch(frame, "opcode", None)
+            if self.async_awaitable is not None:
+                try:
+                    if self.breaks:
+                        with self.set_enterframe(frame):
+                            # set_continue requires enterframe to work
+                            self.set_continue()
+                        self.start_trace()
+                    await self.async_awaitable
+                except Exception:
+                    self._error_exc()
+            else:
+                break
+
+        self.async_shim_frame = None
+
+        # start the trace (the actual command is already set by set_* calls)
+        if self.returnframe is None and self.stoplineno == -1 and not self.breaks:
+            # This means we did a continue without any breakpoints, we should not
+            # start the trace
+            return
+
+        self.start_trace()
+
     def sigint_handler(self, signum, frame):
         if self.allow_kbdint:
             raise KeyboardInterrupt
@@ -782,12 +836,25 @@ class Pdb(bdb.Bdb, cmd.Cmd):
 
         return True
 
-    def default(self, line):
-        if line[:1] == '!': line = line[1:].strip()
-        locals = self.curframe.f_locals
-        globals = self.curframe.f_globals
+    def _exec_await(self, source, globals, locals):
+        """ Run source code that contains await by playing with async shim frame"""
+        # Put the source in an async function
+        source_async = (
+            "async def __pdb_await():\n" +
+            textwrap.indent(source, "    ") + '\n' +
+            "    __pdb_locals.update(locals())"
+        )
+        ns = globals | locals
+        # We use __pdb_locals to do write back
+        ns["__pdb_locals"] = locals
+        exec(source_async, ns)
+        self.async_awaitable = ns["__pdb_await"]()
+
+    def _read_code(self, line):
+        buffer = line
+        is_await_code = False
+        code = None
         try:
-            buffer = line
             if (code := codeop.compile_command(line + '\n', '<stdin>', 'single')) is None:
                 # Multi-line mode
                 with self._enable_multiline_completion():
@@ -800,7 +867,7 @@ class Pdb(bdb.Bdb, cmd.Cmd):
                             except (EOFError, KeyboardInterrupt):
                                 self.lastcmd = ""
                                 print('\n')
-                                return
+                                return None, None, False
                         else:
                             self.stdout.write(continue_prompt)
                             self.stdout.flush()
@@ -809,11 +876,31 @@ class Pdb(bdb.Bdb, cmd.Cmd):
                                 self.lastcmd = ""
                                 self.stdout.write('\n')
                                 self.stdout.flush()
-                                return
+                                return None, None, False
                             else:
                                 line = line.rstrip('\r\n')
                         buffer += '\n' + line
                     self.lastcmd = buffer
+        except SyntaxError as e:
+            # Maybe it's an await expression/statement
+            if (
+                self.async_shim_frame is not None
+                and e.msg == "'await' outside function"
+            ):
+                is_await_code = True
+            else:
+                raise
+
+        return code, buffer, is_await_code
+
+    def default(self, line):
+        if line[:1] == '!': line = line[1:].strip()
+        locals = self.curframe.f_locals
+        globals = self.curframe.f_globals
+        try:
+            code, buffer, is_await_code = self._read_code(line)
+            if buffer is None:
+                return
             save_stdout = sys.stdout
             save_stdin = sys.stdin
             save_displayhook = sys.displayhook
@@ -821,8 +908,12 @@ class Pdb(bdb.Bdb, cmd.Cmd):
                 sys.stdin = self.stdin
                 sys.stdout = self.stdout
                 sys.displayhook = self.displayhook
-                if not self._exec_in_closure(buffer, globals, locals):
-                    exec(code, globals, locals)
+                if is_await_code:
+                    self._exec_await(buffer, globals, locals)
+                    return True
+                else:
+                    if not self._exec_in_closure(buffer, globals, locals):
+                        exec(code, globals, locals)
             finally:
                 sys.stdout = save_stdout
                 sys.stdin = save_stdin
@@ -2501,6 +2592,21 @@ def set_trace(*, header=None, commands=None):
         pdb.message(header)
     pdb.set_trace(sys._getframe().f_back, commands=commands)
 
+async def set_trace_async(*, header=None, commands=None):
+    """Enter the debugger at the calling stack frame, but in async mode.
+
+    This should be used as await pdb.set_trace_async(). Users can do await
+    if they enter the debugger with this function. Otherwise it's the same
+    as set_trace().
+    """
+    if Pdb._last_pdb_instance is not None:
+        pdb = Pdb._last_pdb_instance
+    else:
+        pdb = Pdb(mode='inline', backend='monitoring')
+    if header is not None:
+        pdb.message(header)
+    await pdb.set_trace_async(sys._getframe().f_back, commands=commands)
+
 # Remote PDB
 
 class _PdbServer(Pdb):