gh-120678: pyrepl: Include globals from modules passed with `-i` (GH-120904)

Co-authored-by: Łukasz Langa <lukasz@langa.pl>

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 @@ handle_events_narrow_console = partial(
 )
 
 
+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 io
 import itertools
 import os
 import pathlib
+import re
 import rlcompleter
 import select
 import subprocess
@@ -21,7 +22,8 @@ from .support import (
     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 @@ class TestPyReplOutput(TestCase):
         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 @@ class TestMain(TestCase):
         # 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 @@ class TestMain(TestCase):
 
         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 @@ class TestMain(TestCase):
         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 @@ error:
 }
 
 
+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;
 }