gh-83151: Make closure work on pdb

Lib/pdb.py

@@ -77,10 +77,12 @@
 import code
 import glob
 import token
+import types
 import codeop
 import pprint
 import signal
 import inspect
+import textwrap
 import tokenize
 import traceback
 import linecache
@@ -624,11 +626,96 @@ def _disable_command_completion(self):
  self.completenames = completenames
  return
 
+ def _exec_in_closure(self, source, globals, locals):
+ """ Run source code in closure so code object created within source
+ can find variables in locals correctly
+
+ returns True if the source is executed, False otherwise
+ """
+
+ # Determine if the source should be executed in closure. Only when the
+ # source compiled to multiple code objects, we should use this feature.
+ # Otherwise, we can just raise an exception and normal exec will be used.
+
+ code = compile(source, "<string>", "exec")
+ if not any(isinstance(const, CodeType) for const in code.co_consts):
+ return False
+
+ # locals could be a proxy which does not support pop
+ # copy it first to avoid modifying the original locals
+ locals_copy = dict(locals)
+
+ locals_copy["__pdb_eval__"] = {
+ "result": None,
+ "write_back": {}
+ }
+
+ # If the source is an expression, we need to print its value
+ try:
+ compile(source, "<string>", "eval")
+ except SyntaxError:
+ pass
+ else:
+ source = "__pdb_eval__['result'] = " + source
+
+ # Add write-back to update the locals
+ source = ("try:\n" +
+ textwrap.indent(source, " ") + "\n" +
+ "finally:\n" +
+ " __pdb_eval__['write_back'] = locals()")
+
+ # Build a closure source code with freevars from locals like:
+ # def __pdb_outer():
+ # var = None
+ # def __pdb_scope(): # This is the code object we want to execute
+ # nonlocal var
+ # <source>
+ # return __pdb_scope.__code__
+ source_with_closure = ("def __pdb_outer():\n" +
+ "\n".join(f" {var} = None" for var in locals_copy) + "\n" +
+ " def __pdb_scope():\n" +
+ "\n".join(f" nonlocal {var}" for var in locals_copy) + "\n" +
+ textwrap.indent(source, " ") + "\n" +
+ " return __pdb_scope.__code__"
+ )
+
+ # Get the code object of __pdb_scope()
+ # The exec fills locals_copy with the __pdb_outer() function and we can call
+ # that to get the code object of __pdb_scope()
+ ns = {}
+ try:
+ exec(source_with_closure, {}, ns)
+ except Exception:
+ return False
+ code = ns["__pdb_outer"]()
+
+ cells = tuple(types.CellType(locals_copy.get(var)) for var in code.co_freevars)
+
+ try:
+ exec(code, globals, locals_copy, closure=cells)
+ except Exception:
+ return False
+
+ # get the data we need from the statement
+ pdb_eval = locals_copy["__pdb_eval__"]
+
+ # __pdb_eval__ should not be updated back to locals
+ pdb_eval["write_back"].pop("__pdb_eval__")
+
+ # Write all local variables back to locals
+ locals.update(pdb_eval["write_back"])
+ eval_result = pdb_eval["result"]
+ if eval_result is not None:
+ print(repr(eval_result))
+
+ return True
+
  def default(self, line):
  if line[:1] == '!': line = line[1:].strip()
  locals = self.curframe_locals
  globals = self.curframe.f_globals
  try:
+ buffer = line
  if (code := codeop.compile_command(line + '\n', '<stdin>', 'single')) is None:
  # Multi-line mode
  with self._disable_command_completion():
@@ -661,7 +748,8 @@ def default(self, line):
  sys.stdin = self.stdin
  sys.stdout = self.stdout
  sys.displayhook = self.displayhook
- exec(code, globals, locals)
+ if not self._exec_in_closure(buffer, globals, locals):
+ exec(code, globals, locals)
  finally:
  sys.stdout = save_stdout
  sys.stdin = save_stdin

Lib/test/test_pdb.py

@@ -2224,8 +2224,71 @@ def test_pdb_multiline_statement():
  (Pdb) c
  """
 
+def test_pdb_closure():
+ """Test for all expressions/statements that involve closure
+
+ >>> k = 0
+ >>> g = 1
+ >>> def test_function():
+ ... x = 2
+ ... g = 3
+ ... import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
+
+ >>> with PdbTestInput([ # doctest: +NORMALIZE_WHITESPACE
+ ... 'k',
+ ... 'g',
+ ... 'y = y',
+ ... 'global g; g',
+ ... 'global g; (lambda: g)()',
+ ... '(lambda: x)()',
+ ... '(lambda: g)()',
+ ... 'lst = [n for n in range(10) if (n % x) == 0]',
+ ... 'lst',
+ ... 'sum(n for n in lst if n > x)',
+ ... 'x = 1; raise Exception()',
+ ... 'x',
+ ... 'def f():',
+ ... ' return x',
+ ... '',
+ ... 'f()',
+ ... 'c'
+ ... ]):
+ ... test_function()
+ > <doctest test.test_pdb.test_pdb_closure[2]>(4)test_function()
+ -> import pdb; pdb.Pdb(nosigint=True, readrc=False).set_trace()
+ (Pdb) k
+ 0
+ (Pdb) g
+ 3
+ (Pdb) y = y
+ *** NameError: name 'y' is not defined
+ (Pdb) global g; g
+ 1
+ (Pdb) global g; (lambda: g)()
+ 1
+ (Pdb) (lambda: x)()
+ 2
+ (Pdb) (lambda: g)()
+ 3
+ (Pdb) lst = [n for n in range(10) if (n % x) == 0]
+ (Pdb) lst
+ [0, 2, 4, 6, 8]
+ (Pdb) sum(n for n in lst if n > x)
+ 18
+ (Pdb) x = 1; raise Exception()
+ *** Exception
+ (Pdb) x
+ 1
+ (Pdb) def f():
+ ... return x
+ ...
+ (Pdb) f()
+ 1
+ (Pdb) c
+ """
+
 def test_pdb_show_attribute_and_item():
- """Test for multiline statement
+ """Test for expressions with command prefix
 
  >>> def test_function():
  ... n = lambda x: x

Misc/NEWS.d/next/Library/2023-10-20-03-50-17.gh-issue-83151.bcsD40.rst

@@ -0,0 +1,3 @@
+Enabled arbitrary statements and evaluations in :mod:`pdb` shell to access the
+local variables of the current frame, which made it possible for multi-scope
+code like generators or nested function to work.