diff --git a/Lib/_pyrepl/reader.py b/Lib/_pyrepl/reader.py index 0ebd9162eca4bb..e6514f0e5eb683 100644 --- a/Lib/_pyrepl/reader.py +++ b/Lib/_pyrepl/reader.py @@ -28,6 +28,7 @@ from dataclasses import dataclass, field, fields from . import commands, console, input +from .utils import DEFAULT_PS1, DEFAULT_PS2, DEFAULT_PS3, DEFAULT_PS4 from .utils import wlen, unbracket, disp_str, gen_colors, THEME from .trace import trace @@ -473,22 +474,40 @@ def get_arg(self, default: int = 1) -> int: return default return self.arg + @staticmethod + def __get_prompt_str(prompt: object, default_prompt: str) -> str: + """ + Convert prompt object to string. + + If str(prompt) raises BaseException, MemoryError or 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 = DEFAULT_PS1 + arg = self.__get_prompt_str(self.arg, "") + if arg: + prompt = f"(arg: {self.arg}) " elif self.paste_mode: 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: t = THEME() diff --git a/Lib/_pyrepl/utils.py b/Lib/_pyrepl/utils.py index e04fbdc6c8a5c4..77f5558ab7b237 100644 --- a/Lib/_pyrepl/utils.py +++ b/Lib/_pyrepl/utils.py @@ -58,6 +58,12 @@ class ColorSpan(NamedTuple): tag: str +DEFAULT_PS1 = ">>> " +DEFAULT_PS2 = ">>> " +DEFAULT_PS3 = "... " +DEFAULT_PS4 = "... " + + @functools.cache def str_width(c: str) -> int: if ord(c) < 128: diff --git a/Lib/test/test_pyrepl/test_reader.py b/Lib/test/test_pyrepl/test_reader.py index 1f655264f1c00a..d5c57fbb6c5df1 100644 --- a/Lib/test/test_pyrepl/test_reader.py +++ b/Lib/test/test_pyrepl/test_reader.py @@ -11,9 +11,27 @@ from .support import 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 from _colorize import default_theme +def prepare_reader_with_prompt( + console, ps1=DEFAULT_PS1, ps2=DEFAULT_PS2, ps3=DEFAULT_PS3, ps4=DEFAULT_PS4): + reader = prepare_reader( + console, + can_colorize=False, + paste_mode=False, + ps1=ps1, + ps2=ps2, + ps3=ps3, + ps4=ps4 + ) + + # we should use original get_prompt from reader to get exceptions + del reader.get_prompt + return reader + + overrides = {"reset": "z", "soft_keyword": "K"} colors = {overrides.get(k, k[0].lower()): v for k, v in default_theme.syntax.items()} @@ -300,6 +318,89 @@ 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 exceptions from ps1 prompt + class Prompt: + def __str__(self): 1/0 + + _prepare_reader = functools.partial( + prepare_reader_with_prompt, + ps1=Prompt(), + ) + + reader, _ = handle_all_events( + events=code_to_events("a=1"), + prepare_reader=_prepare_reader + ) + + prompt = reader.get_prompt(0, False) + self.assertEqual(prompt, DEFAULT_PS1) + + def test_prompt_ps2_ps3_ps4_raise_exception(self): + # Handles exceptions from ps2, ps3 and ps4 prompts + class Prompt: + def __str__(self): 1/0 + + _prepare_reader = functools.partial( + prepare_reader_with_prompt, + ps1=Prompt(), + ps2=Prompt(), + ps3=Prompt(), + ps4=Prompt(), + ) + + reader, _ = handle_all_events( + events=code_to_events("if cond:\nfunc()\nfunc()"), + prepare_reader=_prepare_reader + ) + + prompt = reader.get_prompt(0, False) + self.assertEqual(prompt, DEFAULT_PS2) + + prompt = reader.get_prompt(1, False) + self.assertEqual(prompt, DEFAULT_PS3) + + prompt = reader.get_prompt(2, False) + self.assertEqual(prompt, DEFAULT_PS4) + + def test_prompt_arg_raise_exception(self): + # Handles exceptions from arg prompt + class Prompt: + def __str__(self): 1/0 + def __rmul__(self, b): return b + + reader, _ = handle_all_events( + events=code_to_events("if some_condition:\nsome_function()"), + prepare_reader=prepare_reader_with_prompt, + ) + + reader.arg = Prompt() + prompt = reader.get_prompt(0, True) + self.assertEqual(prompt, DEFAULT_PS1) + + def test_prompt_raise_exception(self): + # Tests unrecoverable exceptions from prompts + cases = [ + (MemoryError, "No memory for prompt"), + (SystemError, "System error for prompt"), + ] + for cls, msg in cases: + with self.subTest(msg): + + class Prompt: + def __str__(self): raise cls(msg) + + _prepare_reader = functools.partial( + prepare_reader_with_prompt, + ps1=Prompt(), + ) + + with self.assertRaisesRegex(cls, msg): + handle_events_narrow_console( + events=code_to_events("a=1"), + prepare_reader=_prepare_reader, + ) + def test_completions_updated_on_key_press(self): namespace = {"itertools": itertools} code = "itertools." diff --git a/Misc/NEWS.d/next/Library/2025-03-13-00-39-54.gh-issue-130698.o2aU3e.rst b/Misc/NEWS.d/next/Library/2025-03-13-00-39-54.gh-issue-130698.o2aU3e.rst new file mode 100644 index 00000000000000..609345dddc04fa --- /dev/null +++ b/Misc/NEWS.d/next/Library/2025-03-13-00-39-54.gh-issue-130698.o2aU3e.rst @@ -0,0 +1 @@ +Avoid exiting the new REPL when prompt object raises an exception.