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

This commit is contained in:
Daniel Hollas 2024-06-03 18:07:06 +01:00 committed by GitHub
parent 4e8aa32245
commit 2e0aa731ae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 41 additions and 2 deletions

View file

@ -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)