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

1 files changed, 8 insertions, 2 deletions

diff --git a/Lib/_pyrepl/reader.py b/Lib/_pyrepl/reader.py
index 5401ae7..f2e68ef 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