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

3 files changed, 102 insertions, 29 deletions

diff --git a/Lib/code.py b/Lib/code.py
index 2bd5fa3..f4aecdd 100644
--- a/Lib/code.py
+++ b/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 @@ class InteractiveConsole(InteractiveInterpreter):
         """
         InteractiveInterpreter.__init__(self, locals)
         self.filename = filename
+        self.local_exit = local_exit
         self.resetbuffer()
 
     def resetbuffer(self):
@@ -219,27 +221,64 @@ class InteractiveConsole(InteractiveInterpreter):
         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 @@ class InteractiveConsole(InteractiveInterpreter):
         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:
diff --git a/Lib/pdb.py b/Lib/pdb.py
index 67f8d57..1e4d0a2 100755
--- a/Lib/pdb.py
+++ b/Lib/pdb.py
@@ -1741,7 +1741,7 @@ class Pdb(bdb.Bdb, cmd.Cmd):
         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]]
diff --git a/Lib/test/test_code_module.py b/Lib/test/test_code_module.py
index 226bc3a..747c0f9 100644
--- a/Lib/test/test_code_module.py
+++ b/Lib/test/test_code_module.py
@@ -10,11 +10,7 @@ from test.support import import_helper
 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 @@ class TestInteractiveConsole(unittest.TestCase):
         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 @@ class TestInteractiveConsole(unittest.TestCase):
         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()