Add set_trace_async for async support

python · gaogaotiantian · Apr 29, 2025 · Apr 13, 2025 · Apr 16, 2025 · Apr 16, 2025
commit 1b77cdbde71c0c3925def1d71602cda2070c1341
@@ -378,6 +378,9 @@ def __init__(self, completekey='tab', stdin=None, stdout=None, skip=None,
         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
 
@@ -393,6 +396,57 @@ def set_trace(self, frame=None, *, commands=None):
 
         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
@@ -775,6 +829,20 @@ def _exec_in_closure(self, source, globals, locals):
 
         return True
 
+    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 default(self, line):
         if line[:1] == '!': line = line[1:].strip()
         locals = self.curframe.f_locals
@@ -820,8 +888,20 @@ def default(self, line):
                 sys.stdout = save_stdout
                 sys.stdin = save_stdin
                 sys.displayhook = save_displayhook
-        except:
-            self._error_exc()
+        except Exception as e:
+            # Maybe it's an await expression/statement
+            if (
+                isinstance(e, SyntaxError)
+                and e.msg == "'await' outside function"
+                and self.async_shim_frame is not None
+            ):
+                try:
+                    self._exec_await(buffer, globals, locals)
+                    return True
+                except:
+                    self._error_exc()
+            else:
+                self._error_exc()
 
     def _replace_convenience_variables(self, line):
         """Replace the convenience variables in 'line' with their values.
@@ -2491,6 +2571,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)
+
 # Post-Mortem interface
 
 def post_mortem(t=None):

@@ -2169,6 +2169,70 @@ def test_pdb_asynctask():
             (Pdb) continue
             """
 
+        def test_pdb_await_support():
+            """Testing await support in pdb
+
+            >>> import asyncio
+
+            >>> async def test():
+            ...     print("hello")
+            ...     await asyncio.sleep(0)
+            ...     print("world")
+            ...     return 42
+
+            >>> async def main():
+            ...     import pdb;
+            ...     task = asyncio.create_task(test())
+            ...     await pdb.Pdb(nosigint=True, readrc=False).set_trace_async()
+            ...     pass
+
+            >>> def test_function():
+            ...     asyncio.run(main(), loop_factory=asyncio.EventLoop)
+
+            >>> with PdbTestInput([  # doctest: +ELLIPSIS
+            ...     'x = await task',
+            ...     'p x',
+            ...     'x = await test()',
+            ...     'p x',
+            ...     'new_task = asyncio.create_task(test())',
+            ...     'await new_task',
+            ...     'await non_exist()',
+            ...     's',
+            ...     'continue',
+            ... ]):
+            ...     test_function()
+            > <doctest test.test_pdb.test_pdb_await_support[2]>(4)main()
+            -> await pdb.Pdb(nosigint=True, readrc=False).set_trace_async()
+            (Pdb) x = await task
+            hello
+            world
+            > <doctest test.test_pdb.test_pdb_await_support[2]>(4)main()
+            -> await pdb.Pdb(nosigint=True, readrc=False).set_trace_async()
+            (Pdb) p x
+            42
+            (Pdb) x = await test()
+            hello
+            world
+            > <doctest test.test_pdb.test_pdb_await_support[2]>(4)main()
+            -> await pdb.Pdb(nosigint=True, readrc=False).set_trace_async()
+            (Pdb) p x
+            42
+            (Pdb) new_task = asyncio.create_task(test())
+            (Pdb) await new_task
+            hello
+            world
+            > <doctest test.test_pdb.test_pdb_await_support[2]>(4)main()
+            -> await pdb.Pdb(nosigint=True, readrc=False).set_trace_async()
+            (Pdb) await non_exist()
+            *** NameError: name 'non_exist' is not defined
+            > <doctest test.test_pdb.test_pdb_await_support[2]>(4)main()
+            -> await pdb.Pdb(nosigint=True, readrc=False).set_trace_async()
+            (Pdb) s
+            > <doctest test.test_pdb.test_pdb_await_support[2]>(5)main()
+            -> pass
+            (Pdb) continue
+            """
+
     def test_pdb_next_command_for_coroutine():
         """Testing skip unwinding stack on yield for coroutines for "next" command