From 5a840c0bbd9edc1d9c46999a8821b544e1a918f5 Mon Sep 17 00:00:00 2001
From: Jonathan Slenders <jonathan@slenders.be>
Date: Mon, 22 Jul 2024 09:21:01 +0000
Subject: [PATCH] Add custom 'exit' function to return from REPL.

- This doesn't doesn't terminate sys.stdin when `exit` is called, whis is important for `embed()`.
---
 ptpython/repl.py | 43 ++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 40 insertions(+), 3 deletions(-)

diff --git a/ptpython/repl.py b/ptpython/repl.py
index ea2d84f..b6fc0a5 100644
--- a/ptpython/repl.py
+++ b/ptpython/repl.py
@@ -20,7 +20,7 @@
 import warnings
 from dis import COMPILER_FLAG_NAMES
 from pathlib import Path
-from typing import Any, Callable, ContextManager, Iterable, Sequence
+from typing import Any, Callable, ContextManager, Iterable, NoReturn, Sequence
 
 from prompt_toolkit.formatted_text import OneStyleAndTextTuple
 from prompt_toolkit.patch_stdout import patch_stdout as patch_stdout_context
@@ -40,7 +40,15 @@
 except ImportError:
     PyCF_ALLOW_TOP_LEVEL_AWAIT = 0
 
-__all__ = ["PythonRepl", "enable_deprecation_warnings", "run_config", "embed"]
+
+__all__ = [
+    "PythonRepl",
+    "enable_deprecation_warnings",
+    "run_config",
+    "embed",
+    "exit",
+    "ReplExit",
+]
 
 
 def _get_coroutine_flag() -> int | None:
@@ -91,6 +99,8 @@ def run_and_show_expression(self, expression: str) -> None:
                 raise
             except SystemExit:
                 raise
+            except ReplExit:
+                raise
             except BaseException as e:
                 self._handle_exception(e)
             else:
@@ -155,7 +165,10 @@ def run(self) -> None:
                     continue
 
                 # Run it; display the result (or errors if applicable).
-                self.run_and_show_expression(text)
+                try:
+                    self.run_and_show_expression(text)
+                except ReplExit:
+                    return
         finally:
             if self.terminal_title:
                 clear_title()
@@ -383,6 +396,7 @@ def get_ptpython() -> PythonInput:
             return self
 
         globals["get_ptpython"] = get_ptpython
+        globals["exit"] = exit()
 
     def _remove_from_namespace(self) -> None:
         """
@@ -459,6 +473,29 @@ def enter_to_continue() -> None:
         enter_to_continue()
 
 
+class exit:
+    """
+    Exit the ptpython REPL.
+    """
+
+    # This custom exit function ensures that the `embed` function returns from
+    # where we are embedded, and Python doesn't close `sys.stdin` like
+    # the default `exit` from `_sitebuiltins.Quitter` does.
+
+    def __call__(self) -> NoReturn:
+        raise ReplExit
+
+    def __repr__(self) -> str:
+        # (Same message as the built-in Python REPL.)
+        return "Use exit() or Ctrl-D (i.e. EOF) to exit"
+
+
+class ReplExit(Exception):
+    """
+    Exception raised by ptpython's exit function.
+    """
+
+
 def embed(
     globals: dict[str, Any] | None = None,
     locals: dict[str, Any] | None = None,