gh-118835: pyrepl: Fix prompt length computation for custom prompts containing ANSI escape codes

Lib/_pyrepl/reader.py

@@ -28,7 +28,7 @@
 
 
 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 @@ def calc_complete_screen(self) -> list[str]:
                 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 @@ def process_prompt(self, prompt: str) -> tuple[str, int]:
         # 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

Lib/test/test_pyrepl/test_reader.py

@@ -4,6 +4,7 @@
 
 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 @@ def test_newline_within_block_trailing_whitespace(self):
         )
         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)

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.