gh-130698: Add safe methods to get prompts for new REPL

Author: sergey-miryanov

Lib/_pyrepl/reader.py

@@ -31,6 +31,7 @@
 
 from . import commands, console, input
 from .utils import ANSI_ESCAPE_SEQUENCE, wlen, str_width
+from .utils import DEFAULT_PS1, DEFAULT_PS2, DEFAULT_PS3, DEFAULT_PS4
 from .trace import trace
 
 
@@ -531,22 +532,37 @@ def get_arg(self, default: int = 1) -> int:
             return default
         return self.arg
 
+    @staticmethod
+    def __get_prompt_str(prompt, default_prompt) -> str:
+        """
+        Convert prompt object to string.
+
+        If prompt raise BaseException, MemoryError and SystemError then stop
+        the REPL. For other exceptions return default_prompt.
+        """
+        try:
+            return str(prompt)
+        except (MemoryError, SystemError):
+            raise
+        except Exception:
+            return default_prompt
+
     def get_prompt(self, lineno: int, cursor_on_line: bool) -> str:
         """Return what should be in the left-hand margin for line
         'lineno'."""
         if self.arg is not None and cursor_on_line:
-            prompt = f"(arg: {self.arg}) "
+            prompt = self.__get_prompt_str(f"(arg: {self.arg}) ", DEFAULT_PS1)
         elif self.paste_mode and not self.in_bracketed_paste:
             prompt = "(paste) "
         elif "\n" in self.buffer:
             if lineno == 0:
-                prompt = self.ps2
+                prompt = self.__get_prompt_str(self.ps2, DEFAULT_PS2)
             elif self.ps4 and lineno == self.buffer.count("\n"):
-                prompt = self.ps4
+                prompt = self.__get_prompt_str(self.ps4, DEFAULT_PS4)
             else:
-                prompt = self.ps3
+                prompt = self.__get_prompt_str(self.ps3, DEFAULT_PS3)
         else:
-            prompt = self.ps1
+            prompt = self.__get_prompt_str(self.ps1, DEFAULT_PS1)
 
         if self.can_colorize:
             prompt = f"{ANSIColors.BOLD_MAGENTA}{prompt}{ANSIColors.RESET}"

Lib/_pyrepl/utils.py

@@ -5,6 +5,12 @@
 ANSI_ESCAPE_SEQUENCE = re.compile(r"\x1b\[[ -@]*[A-~]")
 
 
+DEFAULT_PS1 = ">>> "
+DEFAULT_PS2 = "... "
+DEFAULT_PS3 = "... "
+DEFAULT_PS4 = "... "
+
+
 @functools.cache
 def str_width(c: str) -> int:
     if ord(c) < 128:

Lib/test/test_pyrepl/test_reader.py

@@ -7,6 +7,7 @@
 from .support import handle_all_events, handle_events_narrow_console, code_to_events, prepare_reader, prepare_console
 from _pyrepl.console import Event
 from _pyrepl.reader import Reader
+from _pyrepl.utils import DEFAULT_PS1, DEFAULT_PS2, DEFAULT_PS3, DEFAULT_PS4
 
 
 class TestReader(TestCase):
@@ -279,6 +280,106 @@ def test_prompt_length(self):
         self.assertEqual(prompt, "\033[0;32m樂>\033[0m> ")
         self.assertEqual(l, 5)
 
+    def test_prompt_ps1_raise_exception(self):
+        # Handles simple ASCII prompt
+        class Prompt:
+            def __str__(self): 1/0
+
+        def prepare_reader_keep_prompts(*args, **kwargs):
+            reader = prepare_reader(*args, **kwargs)
+            del reader.get_prompt
+            reader.ps1 = Prompt()
+            reader.ps2 = "+++ "
+            reader.ps3 = "... "
+            reader.ps4 = ""
+            reader.can_colorize = False
+            reader.paste_mode = False
+            return reader
+
+        events = code_to_events("a=1")
+        reader, _ = handle_events_narrow_console(
+            events,
+            prepare_reader=prepare_reader_keep_prompts,
+        )
+
+        prompt = reader.get_prompt(0, False)
+        self.assertEqual(prompt, DEFAULT_PS1)
+
+    def test_prompt_ps2_raise_exception(self):
+        # Handles simple ASCII prompt
+        class Prompt:
+            def __str__(self): 1/0
+
+        def prepare_reader_keep_prompts(*args, **kwargs):
+            reader = prepare_reader(*args, **kwargs)
+            del reader.get_prompt
+            reader.ps1 = "+++ "
+            reader.ps2 = Prompt()
+            reader.ps3 = "--- "
+            reader.ps4 = "~~~ "
+            reader.can_colorize = False
+            reader.paste_mode = False
+            return reader
+
+        events = code_to_events("if some_condition:\nsome_function()")
+        reader, _ = handle_events_narrow_console(
+            events,
+            prepare_reader=prepare_reader_keep_prompts,
+        )
+
+        prompt = reader.get_prompt(0, False)
+        self.assertEqual(prompt, DEFAULT_PS2)
+
+    def test_prompt_ps3_raise_exception(self):
+        # Handles simple ASCII prompt
+        class Prompt:
+            def __str__(self): 1/0
+
+        def prepare_reader_keep_prompts(*args, **kwargs):
+            reader = prepare_reader(*args, **kwargs)
+            del reader.get_prompt
+            reader.ps1 = "+++ "
+            reader.ps2 = "--- "
+            reader.ps3 = Prompt()
+            reader.ps4 = ""
+            reader.can_colorize = False
+            reader.paste_mode = False
+            return reader
+
+        events = code_to_events("if some_condition:\nsome_function()")
+        reader, _ = handle_events_narrow_console(
+            events,
+            prepare_reader=prepare_reader_keep_prompts,
+        )
+
+        prompt = reader.get_prompt(1, False)
+        self.assertEqual(prompt, DEFAULT_PS3)
+
+    def test_prompt_ps4_raise_exception(self):
+        # Handles simple ASCII prompt
+        class Prompt:
+            def __str__(self): 1/0
+
+        def prepare_reader_keep_prompts(*args, **kwargs):
+            reader = prepare_reader(*args, **kwargs)
+            del reader.get_prompt
+            reader.ps1 = "+++ "
+            reader.ps2 = "--- "
+            reader.ps3 = "~~~ "
+            reader.ps4 = Prompt()
+            reader.can_colorize = False
+            reader.paste_mode = False
+            return reader
+
+        events = code_to_events("if some_condition:\nsome_function()")
+        reader, _ = handle_events_narrow_console(
+            events,
+            prepare_reader=prepare_reader_keep_prompts,
+        )
+
+        prompt = reader.get_prompt(1, False)
+        self.assertEqual(prompt, DEFAULT_PS4)
+
     def test_completions_updated_on_key_press(self):
         namespace = {"itertools": itertools}
         code = "itertools."