diff options
| -rw-r--r-- | Lib/_pyrepl/reader.py | 10 | ||||
| -rw-r--r-- | Lib/test/test_pyrepl/test_reader.py | 32 | ||||
| -rw-r--r-- | Misc/NEWS.d/next/Library/2024-06-02-15-09-17.gh-issue-118835.KUAuz6.rst | 1 |
3 files changed, 41 insertions, 2 deletions
diff --git a/Lib/_pyrepl/reader.py b/Lib/_pyrepl/reader.py index 53e0954..7839116 100644 --- a/Lib/_pyrepl/reader.py +++ b/Lib/_pyrepl/reader.py @@ -28,7 +28,7 @@ from _colorize import can_colorize, ANSIColors # type: ignore[import-not-found] from . import commands, console, input -from .utils import ANSI_ESCAPE_SEQUENCE, wlen +from .utils import ANSI_ESCAPE_SEQUENCE, wlen, str_width from .trace import trace @@ -339,7 +339,8 @@ class Reader: screeninfo.append((0, [])) return screen - def process_prompt(self, prompt: str) -> tuple[str, int]: + @staticmethod + def process_prompt(prompt: str) -> tuple[str, int]: """Process the prompt. This means calculate the length of the prompt. The character \x01 @@ -351,6 +352,11 @@ class Reader: # sequences if they were not explicitly within \x01...\x02. # They are CSI (or ANSI) sequences ( ESC [ ... LETTER ) + # wlen from utils already excludes ANSI_ESCAPE_SEQUENCE chars, + # which breaks the logic below so we redefine it here. + def wlen(s: str) -> int: + return sum(str_width(i) for i in s) + out_prompt = "" l = wlen(prompt) pos = 0 diff --git a/Lib/test/test_pyrepl/test_reader.py b/Lib/test/test_pyrepl/test_reader.py index c9b03d5..9fb956b 100644 --- a/Lib/test/test_pyrepl/test_reader.py +++ b/Lib/test/test_pyrepl/test_reader.py @@ -4,6 +4,7 @@ from unittest import TestCase from .support import handle_all_events, handle_events_narrow_console, code_to_events, prepare_reader from _pyrepl.console import Event +from _pyrepl.reader import Reader class TestReader(TestCase): @@ -176,3 +177,34 @@ class TestReader(TestCase): ) self.assert_screen_equals(reader, expected) self.assertTrue(reader.finished) + + def test_prompt_length(self): + # Handles simple ASCII prompt + ps1 = ">>> " + prompt, l = Reader.process_prompt(ps1) + self.assertEqual(prompt, ps1) + self.assertEqual(l, 4) + + # Handles ANSI escape sequences + ps1 = "\033[0;32m>>> \033[0m" + prompt, l = Reader.process_prompt(ps1) + self.assertEqual(prompt, "\033[0;32m>>> \033[0m") + self.assertEqual(l, 4) + + # Handles ANSI escape sequences bracketed in \001 .. \002 + ps1 = "\001\033[0;32m\002>>> \001\033[0m\002" + prompt, l = Reader.process_prompt(ps1) + self.assertEqual(prompt, "\033[0;32m>>> \033[0m") + self.assertEqual(l, 4) + + # Handles wide characters in prompt + ps1 = "樂>> " + prompt, l = Reader.process_prompt(ps1) + self.assertEqual(prompt, ps1) + self.assertEqual(l, 5) + + # Handles wide characters AND ANSI sequences together + ps1 = "\001\033[0;32m\002樂>\001\033[0m\002> " + prompt, l = Reader.process_prompt(ps1) + self.assertEqual(prompt, "\033[0;32m樂>\033[0m> ") + self.assertEqual(l, 5) diff --git a/Misc/NEWS.d/next/Library/2024-06-02-15-09-17.gh-issue-118835.KUAuz6.rst b/Misc/NEWS.d/next/Library/2024-06-02-15-09-17.gh-issue-118835.KUAuz6.rst new file mode 100644 index 0000000..ec9ca20 --- /dev/null +++ b/Misc/NEWS.d/next/Library/2024-06-02-15-09-17.gh-issue-118835.KUAuz6.rst @@ -0,0 +1 @@ +Fix _pyrepl crash when using custom prompt with ANSI escape codes. |