gh-118908: Use __main__ for the default PyREPL namespace (#121054)

ambv · web-flow · commit d611c4c8e989 · 2024-06-26T15:01:10.000-04:00
diff --git a/Lib/_pyrepl/__main__.py b/Lib/_pyrepl/__main__.py
@@ -1,51 +1,3 @@
-import os
-import sys
-
-CAN_USE_PYREPL: bool
-if sys.platform != "win32":
-    CAN_USE_PYREPL = True
-else:
-    CAN_USE_PYREPL = sys.getwindowsversion().build >= 10586  # Windows 10 TH2
-
-
-def interactive_console(mainmodule=None, quiet=False, pythonstartup=False):
-    global CAN_USE_PYREPL
-    if not CAN_USE_PYREPL:
-        return sys._baserepl()
-
-    startup_path = os.getenv("PYTHONSTARTUP")
-    if pythonstartup and startup_path:
-        import tokenize
-        with tokenize.open(startup_path) as f:
-            startup_code = compile(f.read(), startup_path, "exec")
-            exec(startup_code)
-
-    # set sys.{ps1,ps2} just before invoking the interactive interpreter. This
-    # mimics what CPython does in pythonrun.c
-    if not hasattr(sys, "ps1"):
-        sys.ps1 = ">>> "
-    if not hasattr(sys, "ps2"):
-        sys.ps2 = "... "
-
-    run_interactive = None
-    try:
-        import errno
-        if not os.isatty(sys.stdin.fileno()):
-            raise OSError(errno.ENOTTY, "tty required", "stdin")
-        from .simple_interact import check
-        if err := check():
-            raise RuntimeError(err)
-        from .simple_interact import run_multiline_interactive_console
-        run_interactive = run_multiline_interactive_console
-    except Exception as e:
-        from .trace import trace
-        msg = f"warning: can't use pyrepl: {e}"
-        trace(msg)
-        print(msg, file=sys.stderr)
-        CAN_USE_PYREPL = False
-    if run_interactive is None:
-        return sys._baserepl()
-    return run_interactive(mainmodule)
-
 if __name__ == "__main__":
-    interactive_console()
+    from .main import interactive_console as __pyrepl_interactive_console
+    __pyrepl_interactive_console()
diff --git a/Lib/_pyrepl/main.py b/Lib/_pyrepl/main.py
@@ -0,0 +1,55 @@
+import os
+import sys
+
+CAN_USE_PYREPL: bool
+if sys.platform != "win32":
+    CAN_USE_PYREPL = True
+else:
+    CAN_USE_PYREPL = sys.getwindowsversion().build >= 10586  # Windows 10 TH2
+
+
+def interactive_console(mainmodule=None, quiet=False, pythonstartup=False):
+    global CAN_USE_PYREPL
+    if not CAN_USE_PYREPL:
+        return sys._baserepl()
+
+    if mainmodule:
+        namespace = mainmodule.__dict__
+    else:
+        import __main__
+        namespace = __main__.__dict__
+        namespace.pop("__pyrepl_interactive_console", None)
+
+    startup_path = os.getenv("PYTHONSTARTUP")
+    if pythonstartup and startup_path:
+        import tokenize
+        with tokenize.open(startup_path) as f:
+            startup_code = compile(f.read(), startup_path, "exec")
+            exec(startup_code, namespace)
+
+    # set sys.{ps1,ps2} just before invoking the interactive interpreter. This
+    # mimics what CPython does in pythonrun.c
+    if not hasattr(sys, "ps1"):
+        sys.ps1 = ">>> "
+    if not hasattr(sys, "ps2"):
+        sys.ps2 = "... "
+
+    run_interactive = None
+    try:
+        import errno
+        if not os.isatty(sys.stdin.fileno()):
+            raise OSError(errno.ENOTTY, "tty required", "stdin")
+        from .simple_interact import check
+        if err := check():
+            raise RuntimeError(err)
+        from .simple_interact import run_multiline_interactive_console
+        run_interactive = run_multiline_interactive_console
+    except Exception as e:
+        from .trace import trace
+        msg = f"warning: can't use pyrepl: {e}"
+        trace(msg)
+        print(msg, file=sys.stderr)
+        CAN_USE_PYREPL = False
+    if run_interactive is None:
+        return sys._baserepl()
+    run_interactive(namespace)
diff --git a/Lib/_pyrepl/simple_interact.py b/Lib/_pyrepl/simple_interact.py
@@ -80,23 +80,13 @@ def _clear_screen():
     "clear": _clear_screen,
 }
 
-DEFAULT_NAMESPACE: dict[str, Any] = {
-    '__name__': '__main__',
-    '__doc__': None,
-    '__package__': None,
-    '__loader__': None,
-    '__spec__': None,
-    '__annotations__': {},
-    '__builtins__': builtins,
-}
 
 def run_multiline_interactive_console(
-    mainmodule: ModuleType | None = None,
+    namespace: dict[str, Any],
     future_flags: int = 0,
     console: code.InteractiveConsole | None = None,
 ) -> None:
     from .readline import _setup
-    namespace = mainmodule.__dict__ if mainmodule else DEFAULT_NAMESPACE
     _setup(namespace)
 
     if console is None:
diff --git a/Lib/test/test_pyrepl/test_pyrepl.py b/Lib/test/test_pyrepl/test_pyrepl.py
@@ -843,15 +843,26 @@ def test_bracketed_paste_single_line(self):
 class TestMain(TestCase):
     @force_not_colorized
     def test_exposed_globals_in_repl(self):
-        expected_output = (
-            "[\'__annotations__\', \'__builtins__\', \'__doc__\', \'__loader__\', "
-            "\'__name__\', \'__package__\', \'__spec__\']"
-        )
+        pre = "['__annotations__', '__builtins__'"
+        post = "'__loader__', '__name__', '__package__', '__spec__']"
         output, exit_code = self.run_repl(["sorted(dir())", "exit"])
-        if "can\'t use pyrepl" in output:
+        if "can't use pyrepl" in output:
             self.skipTest("pyrepl not available")
         self.assertEqual(exit_code, 0)
-        self.assertIn(expected_output, output)
+
+        # if `__main__` is not a file (impossible with pyrepl)
+        case1 = f"{pre}, '__doc__', {post}" in output
+
+        # if `__main__` is an uncached .py file (no .pyc)
+        case2 = f"{pre}, '__doc__', '__file__', {post}" in output
+
+        # if `__main__` is a cached .pyc file and the .py source exists
+        case3 = f"{pre}, '__cached__', '__doc__', '__file__', {post}" in output
+
+        # if `__main__` is a cached .pyc file but there's no .py source file
+        case4 = f"{pre}, '__cached__', '__doc__', {post}" in output
+
+        self.assertTrue(case1 or case2 or case3 or case4, output)
 
     def test_dumb_terminal_exits_cleanly(self):
         env = os.environ.copy()
