gh-120678: pyrepl: Include globals from modules passed with -i

Lib/_pyrepl/__main__.py

@@ -1,3 +1,6 @@
+# Important: don't add things to this module, as they will end up in the REPL's
+# default globals. Use _pyrepl.main instead.
+
 if __name__ == "__main__":
  from .main import interactive_console as __pyrepl_interactive_console
  __pyrepl_interactive_console()

Lib/test/support/__init__.py

@@ -2614,14 +2614,18 @@ def force_not_colorized(func):
  def wrapper(*args, **kwargs):
  import _colorize
  original_fn = _colorize.can_colorize
- variables ={"PYTHON_COLORS": None, "FORCE_COLOR": None}
+ variables: dict[str, str | None] ={
+ "PYTHON_COLORS": None, "FORCE_COLOR": None, "NO_COLOR": None
+ }
  try:
  for key in variables:
  variables[key] = os.environ.pop(key, None)
+ os.environ["NO_COLOR"] = "1"
  _colorize.can_colorize = lambda: False
  return func(*args, **kwargs)
  finally:
  _colorize.can_colorize = original_fn
+ del os.environ["NO_COLOR"]
  for key, value in variables.items():
  if value is not None:
  os.environ[key] = value

Lib/test/test_pyrepl/support.py

@@ -1,3 +1,4 @@
+import os
 from code import InteractiveConsole
 from functools import partial
 from typing import Iterable
@@ -100,8 +101,18 @@ def handle_all_events(
 )
 
 
+def make_clean_env() -> dict[str, str]:
+ clean_env = os.environ.copy()
+ for k in clean_env.copy():
+ if k.startswith("PYTHON"):
+ clean_env.pop(k)
+ clean_env.pop("FORCE_COLOR", None)
+ clean_env.pop("NO_COLOR", None)
+ return clean_env
+
+
 class FakeConsole(Console):
- def __init__(self, events, encoding="utf-8"):
+ def __init__(self, events, encoding="utf-8") -> None:
  self.events = iter(events)
  self.encoding = encoding
  self.screen = []

Lib/test/test_pyrepl/test_pyrepl.py

@@ -2,6 +2,7 @@
 import itertools
 import os
 import pathlib
+import re
 import rlcompleter
 import select
 import subprocess
@@ -21,7 +22,8 @@
  more_lines,
  multiline_input,
  code_to_events,
- clean_screen
+ clean_screen,
+ make_clean_env,
 )
 from _pyrepl.console import Event
 from _pyrepl.readline import ReadlineAlikeReader, ReadlineConfig
@@ -487,6 +489,18 @@ def prepare_reader(self, events):
  reader.can_colorize = False
  return reader
 
+ def test_stdin_is_tty(self):
+ # Used during test log analysis to figure out if a TTY was available.
+ if os.isatty(sys.stdin.fileno()):
+ return
+ self.skipTest("stdin is not a tty")
+
+ def test_stdout_is_tty(self):
+ # Used during test log analysis to figure out if a TTY was available.
+ if os.isatty(sys.stdout.fileno()):
+ return
+ self.skipTest("stdout is not a tty")
+
  def test_basic(self):
  reader = self.prepare_reader(code_to_events("1+1\n"))
 
@@ -888,12 +902,7 @@ def setUp(self):
  # Cleanup from PYTHON* variables to isolate from local
  # user settings, see #121359. Such variables should be
  # added later in test methods to patched os.environ.
- clean_env = os.environ.copy()
- for k in clean_env.copy():
- if k.startswith("PYTHON"):
- clean_env.pop(k)
-
- patcher = patch('os.environ', new=clean_env)
+ patcher = patch('os.environ', new=make_clean_env())
  self.addCleanup(patcher.stop)
  patcher.start()
 
@@ -920,6 +929,84 @@ def test_exposed_globals_in_repl(self):
 
  self.assertTrue(case1 or case2 or case3 or case4, output)
 
+ def _assertMatchOK(
+ self, var: str, expected: str | re.Pattern, actual: str
+ ) -> None:
+ if isinstance(expected, re.Pattern):
+ self.assertTrue(
+ expected.match(actual),
+ f"{var}={actual} does not match{expected.pattern}",
+ )
+ else:
+ self.assertEqual(
+ actual,
+ expected,
+ f"expected{var}={expected}, got{var}={actual}",
+ )
+
+ @force_not_colorized
+ def _run_repl_globals_test(self, expectations, *, as_file=False, as_module=False):
+ clean_env = make_clean_env()
+ clean_env["NO_COLOR"] = "1" # force_not_colorized doesn't touch subprocesses
+
+ with tempfile.TemporaryDirectory() as td:
+ blue = pathlib.Path(td) / "blue"
+ blue.mkdir()
+ mod = blue / "calx.py"
+ mod.write_text("FOO = 42", encoding="utf-8")
+ commands = [
+ "print(f'{" + var + "=}')" for var in expectations
+ ] + ["exit"]
+ if as_file and as_module:
+ self.fail("as_file and as_module are mutually exclusive")
+ elif as_file:
+ output, exit_code = self.run_repl(
+ commands,
+ cmdline_args=[str(mod)],
+ env=clean_env,
+ )
+ elif as_module:
+ output, exit_code = self.run_repl(
+ commands,
+ cmdline_args=["-m", "blue.calx"],
+ env=clean_env,
+ cwd=td,
+ )
+ else:
+ self.fail("Choose one of as_file or as_module")
+
+ if "can't use pyrepl" in output:
+ self.skipTest("pyrepl not available")
+
+ self.assertEqual(exit_code, 0)
+ for var, expected in expectations.items():
+ with self.subTest(var=var, expected=expected):
+ if m := re.search(rf"[\r\n]{var}=(.+?)[\r\n]", output):
+ self._assertMatchOK(var, expected, actual=m.group(1))
+ else:
+ self.fail(f"{var}= not found in output")
+
+ self.assertNotIn("Exception", output)
+ self.assertNotIn("Traceback", output)
+
+ def test_inspect_keeps_globals_from_inspected_file(self):
+ expectations ={
+ "FOO": "42",
+ "__name__": "'__main__'",
+ "__package__": "None",
+ # "__file__" is missing in -i, like in the basic REPL
+ }
+ self._run_repl_globals_test(expectations, as_file=True)
+
+ def test_inspect_keeps_globals_from_inspected_module(self):
+ expectations ={
+ "FOO": "42",
+ "__name__": "'__main__'",
+ "__package__": "'blue'",
+ "__file__": re.compile(r"^'.*calx.py'$"),
+ }
+ self._run_repl_globals_test(expectations, as_module=True)
+
  def test_dumb_terminal_exits_cleanly(self):
  env = os.environ.copy()
  env.update({"TERM": "dumb"})
@@ -981,16 +1068,27 @@ def test_not_wiping_history_file(self):
  self.assertIn("spam", output)
  self.assertNotEqual(pathlib.Path(hfile.name).stat().st_size, 0)
 
- def run_repl(self, repl_input: str | list[str], env: dict | None = None) -> tuple[str, int]:
+ def run_repl(
+ self,
+ repl_input: str | list[str],
+ env: dict | None = None,
+ *,
+ cmdline_args: list[str] | None = None,
+ cwd: str | None = None,
+ ) -> tuple[str, int]:
+ assert pty
  master_fd, slave_fd = pty.openpty()
  cmd = [sys.executable, "-i", "-u"]
  if env is None:
  cmd.append("-I")
+ if cmdline_args is not None:
+ cmd.extend(cmdline_args)
  process = subprocess.Popen(
  cmd,
  stdin=slave_fd,
  stdout=slave_fd,
  stderr=slave_fd,
+ cwd=cwd,
  text=True,
  close_fds=True,
  env=env if env else os.environ,

Misc/NEWS.d/next/Library/2024-06-22-17-01-56.gh-issue-120678.Ik8dCg.rst

@@ -0,0 +1,3 @@
+Fix regression in the new REPL that meant that globals from files passed
+using the ``-i`` argument would not be included in the REPL's global
+namespace. Patch by Alex Waygood.

Modules/main.c

@@ -4,6 +4,7 @@
 #include "pycore_call.h" // _PyObject_CallNoArgs()
 #include "pycore_initconfig.h" // _PyArgv
 #include "pycore_interp.h" // _PyInterpreterState.sysdict
+#include "pycore_long.h" // _PyLong_GetOne()
 #include "pycore_pathconfig.h" // _PyPathConfig_ComputeSysPath0()
 #include "pycore_pylifecycle.h" // _Py_PreInitializeFromPyArgv()
 #include "pycore_pystate.h" // _PyInterpreterState_GET()
@@ -259,6 +260,53 @@ pymain_run_command(wchar_t *command)
 }
 
 
+static int
+pymain_start_pyrepl_no_main(void)
+{
+ int res = 0;
+ PyObject *pyrepl, *console, *empty_tuple, *kwargs, *console_result;
+ pyrepl = PyImport_ImportModule("_pyrepl.main");
+ if (pyrepl == NULL){
+ fprintf(stderr, "Could not import _pyrepl.main\n");
+ res = pymain_exit_err_print();
+ goto done;
+ }
+ console = PyObject_GetAttrString(pyrepl, "interactive_console");
+ if (console == NULL){
+ fprintf(stderr, "Could not access _pyrepl.main.interactive_console\n");
+ res = pymain_exit_err_print();
+ goto done;
+ }
+ empty_tuple = PyTuple_New(0);
+ if (empty_tuple == NULL){
+ res = pymain_exit_err_print();
+ goto done;
+ }
+ kwargs = PyDict_New();
+ if (kwargs == NULL){
+ res = pymain_exit_err_print();
+ goto done;
+ }
+ if (!PyDict_SetItemString(kwargs, "pythonstartup", _PyLong_GetOne())){
+ _PyRuntime.signals.unhandled_keyboard_interrupt = 0;
+ console_result = PyObject_Call(console, empty_tuple, kwargs);
+ if (!console_result && PyErr_Occurred() == PyExc_KeyboardInterrupt){
+ _PyRuntime.signals.unhandled_keyboard_interrupt = 1;
+ }
+ if (console_result == NULL){
+ res = pymain_exit_err_print();
+ }
+ }
+done:
+ Py_XDECREF(console_result);
+ Py_XDECREF(kwargs);
+ Py_XDECREF(empty_tuple);
+ Py_XDECREF(console);
+ Py_XDECREF(pyrepl);
+ return res;
+}
+
+
 static int
 pymain_run_module(const wchar_t *modname, int set_argv0)
 {
@@ -549,7 +597,7 @@ pymain_repl(PyConfig *config, int *exitcode)
  *exitcode = (run != 0);
  return;
  }
- int run = pymain_run_module(L"_pyrepl", 0);
+ int run = pymain_start_pyrepl_no_main();
  *exitcode = (run != 0);
  return;
 }