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

1 files changed, 32 insertions, 0 deletions

diff --git a/Lib/test/test_pyrepl/test_reader.py b/Lib/test/test_pyrepl/test_reader.py
index c9b03d5e711..9fb956b6555 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)