gh-140287: Handle PYTHONSTARTUP script exceptions in the asyncio REPL

Lib/asyncio/__main__.py

@@ -101,7 +101,13 @@ def run(self):
                 import tokenize
                 with tokenize.open(startup_path) as f:
                     startup_code = compile(f.read(), startup_path, "exec")
-                    exec(startup_code, console.locals)
+                    try:
+                        exec(startup_code, console.locals)
+                    # TODO: Revisit in GH-143023
+                    except SystemExit:
+                        raise
+                    except BaseException:
+                        console.showtraceback()
 
             ps1 = getattr(sys, "ps1", ">>> ")
             if CAN_USE_PYREPL:

Lib/test/test_repl.py

@@ -5,6 +5,7 @@
 import subprocess
 import sys
 import unittest
+from contextlib import contextmanager
 from functools import partial
 from textwrap import dedent
 from test import support
@@ -28,7 +29,7 @@
     raise unittest.SkipTest("test module requires subprocess")
 
 
-def spawn_repl(*args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, custom=False, **kw):
+def spawn_repl(*args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, custom=False, isolated=True, **kw):
     """Run the Python REPL with the given arguments.
 
     kw is extra keyword args to pass to subprocess.Popen. Returns a Popen
@@ -42,7 +43,10 @@ def spawn_repl(*args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, custom=F
     # path may be used by PyConfig_Get("module_search_paths") to build the
     # default module search path.
     stdin_fname = os.path.join(os.path.dirname(sys.executable), "<stdin>")
-    cmd_line = [stdin_fname, '-I']
+    cmd_line = [stdin_fname]
+    # Isolated mode implies -EPs and ignores PYTHON* variables.
+    if isolated:
+        cmd_line.append('-I')
     # Don't re-run the built-in REPL from interactive mode
     # if we're testing a custom REPL (such as the asyncio REPL).
     if not custom:
@@ -64,6 +68,16 @@ def spawn_repl(*args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, custom=F
 spawn_asyncio_repl = partial(spawn_repl, "-m", "asyncio", custom=True)
 
 
+@contextmanager
+def new_startup_env(*, code: str, histfile: str = ".pythonhist"):
+    """Create environment variables for a PYTHONSTARTUP script in a temporary directory."""
+    with os_helper.temp_dir() as tmpdir:
+        filename = os.path.join(tmpdir, "pythonstartup.py")
+        with open(filename, "w") as f:
+            f.write('\n'.join(code.splitlines()))
+        yield {"PYTHONSTARTUP": filename, "PYTHON_HISTORY": os.path.join(tmpdir, histfile)}
+
+
 def run_on_interactive_mode(source):
     """Spawn a new Python interpreter, pass the given
     input source code from the stdin and return the
@@ -197,68 +211,6 @@ def foo(x):
         ]
         self.assertEqual(traceback_lines, expected_lines)
 
-    def test_pythonstartup_error_reporting(self):
-        # errors based on https://github.com/python/cpython/issues/137576
-
-        def make_repl(env):
-            return subprocess.Popen(
-                [os.path.join(os.path.dirname(sys.executable), '<stdin>'), "-i"],
-                executable=sys.executable,
-                text=True,
-                stdin=subprocess.PIPE,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.STDOUT,
-                env=env,
-            )
-
-        # case 1: error in user input, but PYTHONSTARTUP is fine
-        with os_helper.temp_dir() as tmpdir:
-            script = os.path.join(tmpdir, "pythonstartup.py")
-            with open(script, "w") as f:
-                f.write("print('from pythonstartup')" + os.linesep)
-
-            env = os.environ.copy()
-            env['PYTHONSTARTUP'] = script
-            env["PYTHON_HISTORY"] = os.path.join(tmpdir, ".pythonhist")
-            p = make_repl(env)
-            p.stdin.write("1/0")
-            output = kill_python(p)
-        expected = dedent("""
-            Traceback (most recent call last):
-              File "<stdin>", line 1, in <module>
-                1/0
-                ~^~
-            ZeroDivisionError: division by zero
-        """)
-        self.assertIn("from pythonstartup", output)
-        self.assertIn(expected, output)
-
-        # case 2: error in PYTHONSTARTUP triggered by user input
-        with os_helper.temp_dir() as tmpdir:
-            script = os.path.join(tmpdir, "pythonstartup.py")
-            with open(script, "w") as f:
-                f.write("def foo():\n    1/0\n")
-
-            env = os.environ.copy()
-            env['PYTHONSTARTUP'] = script
-            env["PYTHON_HISTORY"] = os.path.join(tmpdir, ".pythonhist")
-            p = make_repl(env)
-            p.stdin.write('foo()')
-            output = kill_python(p)
-        expected = dedent("""
-            Traceback (most recent call last):
-              File "<stdin>", line 1, in <module>
-                foo()
-                ~~~^^
-              File "%s", line 2, in foo
-                1/0
-                ~^~
-            ZeroDivisionError: division by zero
-        """) % script
-        self.assertIn(expected, output)
-
-
-
     def test_runsource_show_syntax_error_location(self):
         user_input = dedent("""def f(x, x): ...
                             """)
@@ -292,23 +244,46 @@ def bar(x):
         expected = "(30, None, [\'def foo(x):\\n\', \'    return x + 1\\n\', \'\\n\'], \'<stdin>\')"
         self.assertIn(expected, output, expected)
 
-    def test_asyncio_repl_reaches_python_startup_script(self):
-        with os_helper.temp_dir() as tmpdir:
-            script = os.path.join(tmpdir, "pythonstartup.py")
-            with open(script, "w") as f:
-                f.write("print('pythonstartup done!')" + os.linesep)
-                f.write("exit(0)" + os.linesep)
-
-            env = os.environ.copy()
-            env["PYTHON_HISTORY"] = os.path.join(tmpdir, ".asyncio_history")
-            env["PYTHONSTARTUP"] = script
-            subprocess.check_call(
-                [sys.executable, "-m", "asyncio"],
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                env=env,
-                timeout=SHORT_TIMEOUT,
-            )
+    def test_pythonstartup_success(self):
+        # errors based on https://github.com/python/cpython/issues/137576
+        # case 1: error in user input, but PYTHONSTARTUP is fine
+        startup_code = "print('notice from pythonstartup')"
+        startup_env = self.enterContext(new_startup_env(code=startup_code))
+        # -q to suppress noise
+        p = spawn_repl("-q", env=os.environ | startup_env, isolated=False)
+        p.stdin.write("1/0")
+        output_lines = kill_python(p).splitlines()
+        self.assertEqual(output_lines[0], 'notice from pythonstartup')
+
+        traceback_lines = output_lines[2:-1]
+        expected_lines = [
+            'Traceback (most recent call last):',
+            '  File "<stdin>", line 1, in <module>',
+            '    1/0',
+            '    ~^~',
+            'ZeroDivisionError: division by zero',
+        ]
+        self.assertEqual(traceback_lines, expected_lines)
+
+    def test_pythonstartup_failure(self):
+        # case 2: error in PYTHONSTARTUP triggered by user input
+        startup_code = "def foo():\n    1/0\n"
+        startup_env = self.enterContext(new_startup_env(code=startup_code))
+        # -q to suppress noise
+        p = spawn_repl("-q", env=os.environ | startup_env, isolated=False)
+        p.stdin.write("foo()")
+        traceback_lines = kill_python(p).splitlines()[1:-1]
+        expected_lines = [
+            'Traceback (most recent call last):',
+            '  File "<stdin>", line 1, in <module>',
+            '    foo()',
+            '    ~~~^^',
+            f'  File "{startup_env['PYTHONSTARTUP']}", line 2, in foo',
+            '    1/0',
+            '    ~^~',
+            'ZeroDivisionError: division by zero',
+        ]
+        self.assertEqual(traceback_lines, expected_lines)
 
     @unittest.skipUnless(pty, "requires pty")
     def test_asyncio_repl_is_ok(self):
@@ -365,6 +340,7 @@ def f():
         self.assertEqual(traceback_lines, expected_lines)
 
 
+@support.force_not_colorized_test_class
 class TestAsyncioREPL(unittest.TestCase):
     def test_multiple_statements_fail_early(self):
         user_input = "1 / 0; print(f'afterwards: {1+1}')"
@@ -409,6 +385,55 @@ def test_toplevel_contextvars_async(self):
         expected = "toplevel contextvar test: ok"
         self.assertIn(expected, output, expected)
 
+    def test_pythonstartup_success(self):
+        startup_code = "import sys\nprint('notice from pythonstartup in asyncio repl', file=sys.stderr)"
+        startup_env = self.enterContext(new_startup_env(code=startup_code, histfile=".asyncio_history"))
+        p = spawn_asyncio_repl(env=os.environ | startup_env, stderr=subprocess.PIPE, isolated=False)
+        p.stdin.write("1/0")
+        kill_python(p)
+        output_lines = p.stderr.read().splitlines()
+        p.stderr.close()
+
+        self.assertEqual(output_lines[3], 'notice from pythonstartup in asyncio repl')
+
+        tb_start_lines = output_lines[5:6]
+        tb_final_lines = output_lines[13:]
+        expected_lines = [
+            'Traceback (most recent call last):',
+            '  File "<stdin>", line 1, in <module>',
+            '    1/0',
+            '    ~^~',
+            'ZeroDivisionError: division by zero',
+            '',
+            'exiting asyncio REPL...',
+        ]
+        self.assertEqual(tb_start_lines + tb_final_lines, expected_lines)
+
+    def test_pythonstartup_failure(self):
+        startup_code = "def foo():\n    1/0\n"
+        startup_env = self.enterContext(new_startup_env(code=startup_code, histfile=".asyncio_history"))
+        p = spawn_asyncio_repl(env=os.environ | startup_env, stderr=subprocess.PIPE, isolated=False)
+        p.stdin.write("foo()")
+        kill_python(p)
+        output_lines = p.stderr.read().splitlines()
+        p.stderr.close()
+
+        tb_start_lines = output_lines[4:5]
+        tb_final_lines = output_lines[12:]
+        expected_lines = [
+            'Traceback (most recent call last):',
+            '  File "<stdin>", line 1, in <module>',
+            '    foo()',
+            '    ~~~^^',
+            f'  File "{startup_env['PYTHONSTARTUP']}", line 2, in foo',
+            '    1/0',
+            '    ~^~',
+            'ZeroDivisionError: division by zero',
+            '',
+            'exiting asyncio REPL...',
+        ]
+        self.assertEqual(tb_start_lines + tb_final_lines, expected_lines)
+
 
 if __name__ == "__main__":
     unittest.main()

Misc/NEWS.d/next/Library/2025-10-18-12-13-39.gh-issue-140287.49iU-4.rst

@@ -0,0 +1,2 @@
+The asyncio REPL now properly handles exceptions in ``PYTHONSTARTUP``
+scripts. Patch by Bartosz Sławecki in :gh:`140287`.