GH-102895 Add an option local_exit in code.interact to block exit() from terminating the whole process (GH-102896)

Author: gaogaotiantian

Doc/library/code.rst

@@ -23,27 +23,34 @@ build applications which provide an interactive interpreter prompt.
    ``'__doc__'`` set to ``None``.
 
 
-.. class:: InteractiveConsole(locals=None, filename="<console>")
+.. class:: InteractiveConsole(locals=None, filename="<console>", local_exit=False)
 
    Closely emulate the behavior of the interactive Python interpreter. This class
    builds on :class:`InteractiveInterpreter` and adds prompting using the familiar
-   ``sys.ps1`` and ``sys.ps2``, and input buffering.
+   ``sys.ps1`` and ``sys.ps2``, and input buffering. If *local_exit* is True,
+   ``exit()`` and ``quit()`` in the console will not raise :exc:`SystemExit`, but
+   instead return to the calling code.
 
+   .. versionchanged:: 3.13
+      Added *local_exit* parameter.
 
-.. function:: interact(banner=None, readfunc=None, local=None, exitmsg=None)
+.. function:: interact(banner=None, readfunc=None, local=None, exitmsg=None, local_exit=False)
 
    Convenience function to run a read-eval-print loop.  This creates a new
    instance of :class:`InteractiveConsole` and sets *readfunc* to be used as
    the :meth:`InteractiveConsole.raw_input` method, if provided.  If *local* is
    provided, it is passed to the :class:`InteractiveConsole` constructor for
-   use as the default namespace for the interpreter loop.  The :meth:`interact`
+   use as the default namespace for the interpreter loop.  If *local_exit* is provided,
+   it is passed to the :class:`InteractiveConsole` constructor.  The :meth:`interact`
    method of the instance is then run with *banner* and *exitmsg* passed as the
    banner and exit message to use, if provided.  The console object is discarded
    after use.
 
    .. versionchanged:: 3.6
       Added *exitmsg* parameter.
 
+   .. versionchanged:: 3.13
+      Added *local_exit* parameter.
 
 .. function:: compile_command(source, filename="<input>", symbol="single")
 

Lib/code.py

@@ -5,6 +5,7 @@
 # Inspired by similar code by Jeff Epler and Fredrik Lundh.
 
 
+import builtins
 import sys
 import traceback
 from codeop import CommandCompiler, compile_command
@@ -169,7 +170,7 @@ class InteractiveConsole(InteractiveInterpreter):
 
     """
 
-    def __init__(self, locals=None, filename="<console>"):
+    def __init__(self, locals=None, filename="<console>", local_exit=False):
         """Constructor.
 
         The optional locals argument will be passed to the
@@ -181,6 +182,7 @@ def __init__(self, locals=None, filename="<console>"):
         """
         InteractiveInterpreter.__init__(self, locals)
         self.filename = filename
+        self.local_exit = local_exit
         self.resetbuffer()
 
     def resetbuffer(self):
@@ -219,27 +221,64 @@ def interact(self, banner=None, exitmsg=None):
         elif banner:
             self.write("%s\n" % str(banner))
         more = 0
-        while 1:
-            try:
-                if more:
-                    prompt = sys.ps2
-                else:
-                    prompt = sys.ps1
+
+        # When the user uses exit() or quit() in their interactive shell
+        # they probably just want to exit the created shell, not the whole
+        # process. exit and quit in builtins closes sys.stdin which makes
+        # it super difficult to restore
+        #
+        # When self.local_exit is True, we overwrite the builtins so
+        # exit() and quit() only raises SystemExit and we can catch that
+        # to only exit the interactive shell
+
+        _exit = None
+        _quit = None
+
+        if self.local_exit:
+            if hasattr(builtins, "exit"):
+                _exit = builtins.exit
+                builtins.exit = Quitter("exit")
+
+            if hasattr(builtins, "quit"):
+                _quit = builtins.quit
+                builtins.quit = Quitter("quit")
+
+        try:
+            while True:
                 try:
-                    line = self.raw_input(prompt)
-                except EOFError:
-                    self.write("\n")
-                    break
-                else:
-                    more = self.push(line)
-            except KeyboardInterrupt:
-                self.write("\nKeyboardInterrupt\n")
-                self.resetbuffer()
-                more = 0
-        if exitmsg is None:
-            self.write('now exiting %s...\n' % self.__class__.__name__)
-        elif exitmsg != '':
-            self.write('%s\n' % exitmsg)
+                    if more:
+                        prompt = sys.ps2
+                    else:
+                        prompt = sys.ps1
+                    try:
+                        line = self.raw_input(prompt)
+                    except EOFError:
+                        self.write("\n")
+                        break
+                    else:
+                        more = self.push(line)
+                except KeyboardInterrupt:
+                    self.write("\nKeyboardInterrupt\n")
+                    self.resetbuffer()
+                    more = 0
+                except SystemExit as e:
+                    if self.local_exit:
+                        self.write("\n")
+                        break
+                    else:
+                        raise e
+        finally:
+            # restore exit and quit in builtins if they were modified
+            if _exit is not None:
+                builtins.exit = _exit
+
+            if _quit is not None:
+                builtins.quit = _quit
+
+            if exitmsg is None:
+                self.write('now exiting %s...\n' % self.__class__.__name__)
+            elif exitmsg != '':
+                self.write('%s\n' % exitmsg)
 
     def push(self, line):
         """Push a line to the interpreter.
@@ -276,8 +315,22 @@ def raw_input(self, prompt=""):
         return input(prompt)
 
 
+class Quitter:
+    def __init__(self, name):
+        self.name = name
+        if sys.platform == "win32":
+            self.eof = 'Ctrl-Z plus Return'
+        else:
+            self.eof = 'Ctrl-D (i.e. EOF)'
+
+    def __repr__(self):
+        return f'Use {self.name} or {self.eof} to exit'
+
+    def __call__(self, code=None):
+        raise SystemExit(code)
+
 
-def interact(banner=None, readfunc=None, local=None, exitmsg=None):
+def interact(banner=None, readfunc=None, local=None, exitmsg=None, local_exit=False):
     """Closely emulate the interactive Python interpreter.
 
     This is a backwards compatible interface to the InteractiveConsole
@@ -290,9 +343,10 @@ def interact(banner=None, readfunc=None, local=None, exitmsg=None):
     readfunc -- if not None, replaces InteractiveConsole.raw_input()
     local -- passed to InteractiveInterpreter.__init__()
     exitmsg -- passed to InteractiveConsole.interact()
+    local_exit -- passed to InteractiveConsole.__init__()
 
     """
-    console = InteractiveConsole(local)
+    console = InteractiveConsole(local, local_exit=local_exit)
     if readfunc is not None:
         console.raw_input = readfunc
     else:

Lib/pdb.py

@@ -1741,7 +1741,7 @@ def do_interact(self, arg):
         contains all the (global and local) names found in the current scope.
         """
         ns = {**self.curframe.f_globals, **self.curframe_locals}
-        code.interact("*interactive*", local=ns)
+        code.interact("*interactive*", local=ns, local_exit=True)
 
     def do_alias(self, arg):
         """alias [name [command]]

Lib/test/test_code_module.py

@@ -10,11 +10,7 @@
 code = import_helper.import_module('code')
 
 
-class TestInteractiveConsole(unittest.TestCase):
-
-    def setUp(self):
-        self.console = code.InteractiveConsole()
-        self.mock_sys()
+class MockSys:
 
     def mock_sys(self):
         "Mock system environment for InteractiveConsole"
@@ -32,6 +28,13 @@ def mock_sys(self):
         del self.sysmod.ps1
         del self.sysmod.ps2
 
+
+class TestInteractiveConsole(unittest.TestCase, MockSys):
+
+    def setUp(self):
+        self.console = code.InteractiveConsole()
+        self.mock_sys()
+
     def test_ps1(self):
         self.infunc.side_effect = EOFError('Finished')
         self.console.interact()
@@ -151,5 +154,21 @@ def test_context_tb(self):
         self.assertIn(expected, output)
 
 
+class TestInteractiveConsoleLocalExit(unittest.TestCase, MockSys):
+
+    def setUp(self):
+        self.console = code.InteractiveConsole(local_exit=True)
+        self.mock_sys()
+
+    def test_exit(self):
+        # default exit message
+        self.infunc.side_effect = ["exit()"]
+        self.console.interact(banner='')
+        self.assertEqual(len(self.stderr.method_calls), 2)
+        err_msg = self.stderr.method_calls[1]
+        expected = 'now exiting InteractiveConsole...\n'
+        self.assertEqual(err_msg, ['write', (expected,), {}])
+
+
 if __name__ == "__main__":
     unittest.main()

Misc/NEWS.d/next/Library/2023-03-22-02-01-30.gh-issue-102895.HiEqaZ.rst

@@ -0,0 +1 @@
+Added a parameter ``local_exit`` for :func:`code.interact` to prevent ``exit()`` and ``quit`` from closing ``sys.stdin`` and raise ``SystemExit``.