From 6b10467fbc0b67bf217ea27956b545103c4a0ad5 Mon Sep 17 00:00:00 2001
From: "Miss Islington (bot)"
 <31488909+miss-islington@users.noreply.github.com>
Date: Mon, 3 Jun 2024 19:31:26 +0200
Subject: [3.13] gh-118835: pyrepl: Fix prompt length computation for custom
 prompts containing ANSI escape codes (GH-119942) (#119990)

gh-118835: pyrepl: Fix prompt length computation for custom prompts containing ANSI escape codes (GH-119942)
(cherry picked from commit 2e0aa731aebb8ef3d89ada82f5d39b1bbac65d1f)

Co-authored-by: Daniel Hollas <daniel.hollas@bristol.ac.uk>
---
 Lib/_pyrepl/reader.py                              | 10 +++++--
 Lib/test/test_pyrepl/test_reader.py                | 32 ++++++++++++++++++++++
 .../2024-06-02-15-09-17.gh-issue-118835.KUAuz6.rst |  1 +
 3 files changed, 41 insertions(+), 2 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2024-06-02-15-09-17.gh-issue-118835.KUAuz6.rst

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.
-- 
cgit v0.12