Various fixes to handle wide characters correctly

Many of the fixes incorporated from https://github.com/pypy/pyrepl/pull/34/files
ambv · May 5, 2024 · Oct 20, 2023 · Apr 27, 2024 · Apr 27, 2024 · Apr 27, 2024
commit ba26254debf4b303fc9bf6ddce497b9d2f6e9cd3
@@ -29,7 +29,7 @@
 
 from . import commands, console, input
 
-_r_csi_seq = re.compile(r"\033\[[ -@]*[A-~]")
+ANSI_ESCAPE_SEQUENCE = re.compile(r"\x1b\[[ -@]*[A-~]")
 
 
 # types
@@ -45,6 +45,13 @@ def str_width(c: str) -> int:
     return 2
 
 
+def wlen(s: str) -> int:
+    length = sum(str_width(i) for i in s)
+
+    # remove lengths of any escape sequences
+    return length - sum(len(i) for i in ANSI_ESCAPE_SEQUENCE.findall(s))
+
+
 def disp_str(buffer: str) -> tuple[str, list[int]]:
     """disp_str(buffer:string) -> (string, [int])
 
@@ -55,19 +62,15 @@ def disp_str(buffer: str) -> tuple[str, list[int]]:
     >>> disp_str(chr(3))
     ('^C', [1, 0])
 
-    the list always contains 0s or 1s at present; it could conceivably
-    go higher as and when unicode support happens."""
+    """
     b: list[int] = []
     s: list[str] = []
     for c in buffer:
         if unicodedata.category(c).startswith("C"):
             c = r"\u%04x" % ord(c)
-            b.append(1)
-            b.extend([0] * (len(c) - 1))
-        else:
-            b.append(1)
-            b.extend([0] * (str_width(c) - 1))
         s.append(c)
+        b.append(wlen(c))
+        b.extend([0] * (len(c) - 1))
     return "".join(s), b
 
 
@@ -253,7 +256,7 @@ def __post_init__(self) -> None:
             self.keymap, invalid_cls="invalid-key", character_cls="self-insert"
         )
         self.screeninfo = [(0, [0])]
-        self.cxy = self.pos2xy(self.pos)
+        self.cxy = self.pos2xy()
         self.lxy = (self.pos, 0)
 
     def collect_keymap(self):
@@ -268,40 +271,47 @@ def calc_screen(self) -> list[str]:
         lines = self.get_unicode().split("\n")
         screen: list[str] = []
         screeninfo: list[tuple[int, list[int]]] = []
-        w = self.console.width - 1
-        p = self.pos
+        width = self.console.width - 1
+        pos = self.pos
         for ln, line in enumerate(lines):
             ll = len(line)
-            if 0 <= p <= ll:
+            if 0 <= pos <= ll:
                 if self.msg and not self.msg_at_bottom:
                     for mline in self.msg.split("\n"):
                         screen.append(mline)
                         screeninfo.append((0, []))
-                self.lxy = p, ln
-            prompt = self.get_prompt(ln, ll >= p >= 0)
+                self.lxy = pos, ln
+            prompt = self.get_prompt(ln, ll >= pos >= 0)
             while "\n" in prompt:
                 pre_prompt, _, prompt = prompt.partition("\n")
                 screen.append(pre_prompt)
                 screeninfo.append((0, []))
-            p -= ll + 1
+            pos -= ll + 1
             prompt, lp = self.process_prompt(prompt)
-            if _can_colorize():
-                prompt = _ANSIColors.BOLD_MAGENTA + prompt + _ANSIColors.RESET
             l, l2 = disp_str(line)
-            wrapcount = (len(l) + lp) // w
+            wrapcount = (wlen(l) + lp) // width
             if wrapcount == 0:
                 screen.append(prompt + l)
-                screeninfo.append((lp, l2 + [1]))
+                screeninfo.append((lp, l2))
             else:
-                screen.append(prompt + l[: w - lp] + "\\")
-                screeninfo.append((lp, l2[: w - lp]))
-                for i in range(-lp + w, -lp + wrapcount * w, w):
-                    screen.append(l[i : i + w] + "\\")
-                    screeninfo.append((0, l2[i : i + w]))
-                screen.append(l[wrapcount * w - lp :])
-                screeninfo.append((0, l2[wrapcount * w - lp :] + [1]))
+                for i in range(wrapcount + 1):
+                    prelen = lp if i == 0 else 0
+                    index_to_wrap_before = 0
+                    column = 0
+                    for character_width in l2:
+                        if column + character_width >= width - prelen:
+                            break
+                        index_to_wrap_before += 1
+                        column += character_width
+                    pre = prompt if i == 0 else ""
+                    post = "\\" if i != wrapcount else ""
+                    after = [1] if i != wrapcount else []
+                    screen.append(pre + l[:index_to_wrap_before] + post)
+                    screeninfo.append((prelen, l2[:index_to_wrap_before] + after))
+                    l = l[index_to_wrap_before:]
+                    l2 = l2[index_to_wrap_before:]
         self.screeninfo = screeninfo
-        self.cxy = self.pos2xy(self.pos)
+        self.cxy = self.pos2xy()
         if self.msg and self.msg_at_bottom:
             for mline in self.msg.split("\n"):
                 screen.append(mline)
@@ -321,7 +331,7 @@ def process_prompt(self, prompt: str) -> tuple[str, int]:
         # They are CSI (or ANSI) sequences  ( ESC [ ... LETTER )
 
         out_prompt = ""
-        l = len(prompt)
+        l = wlen(prompt)
         pos = 0
         while True:
             s = prompt.find("\x01", pos)
@@ -333,11 +343,11 @@ def process_prompt(self, prompt: str) -> tuple[str, int]:
             # Found start and end brackets, subtract from string length
             l = l - (e - s + 1)
             keep = prompt[pos:s]
-            l -= sum(map(len, _r_csi_seq.findall(keep)))
+            l -= sum(map(wlen, ANSI_ESCAPE_SEQUENCE.findall(keep)))
             out_prompt += keep + prompt[s + 1 : e]
             pos = e + 1
         keep = prompt[pos:]
-        l -= sum(map(len, _r_csi_seq.findall(keep)))
+        l -= sum(map(wlen, ANSI_ESCAPE_SEQUENCE.findall(keep)))
         out_prompt += keep
         return out_prompt, l
 
@@ -412,20 +422,22 @@ def get_prompt(self, lineno, cursor_on_line) -> str:
         """Return what should be in the left-hand margin for line
         `lineno'."""
         if self.arg is not None and cursor_on_line:
-            return "(arg: %s) " % self.arg
-
-        if self.paste_mode:
-            return '(paste) '
-
-        if "\n" in self.buffer:
+            prompt = "(arg: %s) " % self.arg
+        elif self.paste_mode:
+            prompt = '(paste) '
+        elif "\n" in self.buffer:
             if lineno == 0:
-                return self.ps2
+                prompt = self.ps2
             elif lineno == self.buffer.count("\n"):
-                return self.ps4
+                prompt = self.ps4
             else:
-                return self.ps3
+                prompt = self.ps3
+        else:
+            prompt = self.ps1
 
-        return self.ps1
+        if _can_colorize():
+            prompt = f"{_ANSIColors.BOLD_MAGENTA}{prompt}{_ANSIColors.RESET}"
+        return prompt
 
     def push_input_trans(self, itrans) -> None:
         self.input_trans_stack.append(self.input_trans)
@@ -434,31 +446,25 @@ def push_input_trans(self, itrans) -> None:
     def pop_input_trans(self) -> None:
         self.input_trans = self.input_trans_stack.pop()
 
-    def pos2xy(self, pos: int) -> tuple[int, int]:
+    def pos2xy(self) -> tuple[int, int]:
         """Return the x, y coordinates of position 'pos'."""
         # this *is* incomprehensible, yes.
         y = 0
+        pos = self.pos
         assert 0 <= pos <= len(self.buffer)
         if pos == len(self.buffer):
             y = len(self.screeninfo) - 1
             p, l2 = self.screeninfo[y]
-            return p + len(l2) - 1, y
-        else:
-            for p, l2 in self.screeninfo:
-                l = l2.count(1)
-                if l > pos:
-                    break
-                else:
-                    pos -= l
-                    y += 1
-            c = 0
-            i = 0
-            while c < pos:
-                c += l2[i]
-                i += 1
-            while l2[i] == 0:
-                i += 1
-            return p + i, y
+            return p + sum(l2) + l2.count(0), y
+
+        for p, l2 in self.screeninfo:
+            l = len(l2) - l2.count(0)
+            if l >= pos:
+                break
+            else:
+                pos -= l + 1
+                y += 1
+        return p + sum(l2[:pos]), y
 
     def insert(self, text) -> None:
         """Insert 'text' at the insertion point."""
@@ -468,7 +474,7 @@ def insert(self, text) -> None:
 
     def update_cursor(self) -> None:
         """Move the cursor to reflect changes in self.pos"""
-        self.cxy = self.pos2xy(self.pos)
+        self.cxy = self.pos2xy()
         self.console.move_cursor(*self.cxy)
 
     def after_command(self, cmd: Command) -> None:
@@ -491,10 +497,6 @@ def prepare(self) -> None:
             self.dirty = True
             self.last_command = None
             self.calc_screen()
-            # self.screeninfo = [(0, [0])]
-            # self.cxy = self.pos2xy(self.pos)
-            # self.lxy = (self.pos, 0)
-            # XXX should kill ring be here?
         except BaseException:
             self.restore()
             raise

@@ -33,6 +33,7 @@
 from . import curses
 from .console import Console, Event
 from .fancy_termios import tcgetattr, tcsetattr
+from .reader import wlen
 from .trace import trace
 from .unix_eventqueue import EventQueue
 
@@ -555,56 +556,76 @@ def __setup_movement(self):
 
         self.__move = self.__move_short
 
-    def __write_changed_line(self, y, oldline, newline, px):
+    def __write_changed_line(self, y, oldline, newline, px_coord):
         # this is frustrating; there's no reason to test (say)
         # self.dch1 inside the loop -- but alternative ways of
         # structuring this function are equally painful (I'm trying to
         # avoid writing code generators these days...)
-        x = 0
-        minlen = min(len(oldline), len(newline))
-        #
+        minlen = min(wlen(oldline), wlen(newline))
+        x_pos = 0
+        x_coord = 0
+
+        px_pos = 0
+        j = 0
+        for c in oldline:
+            if j >= px_coord: break
+            j += wlen(c)
+            px_pos += 1
+
         # reuse the oldline as much as possible, but stop as soon as we
         # encounter an ESCAPE, because it might be the start of an escape
         # sequene
-        while x < minlen and oldline[x] == newline[x] and newline[x] != "\x1b":
-            x += 1
-        if oldline[x:] == newline[x + 1 :] and self.ich1:
+        while x_coord < minlen and oldline[x_pos] == newline[x_pos] and newline[x_pos] != "\x1b":
+            x_coord += wlen(newline[x_pos])
+            x_pos += 1
+
+        character_width = wlen(newline[x_pos])
+
+        # if we need to insert a single character right after the first detected change
+        if oldline[x_pos:] == newline[x_pos + 1 :] and self.ich1:
             if (
                 y == self.__posxy[1]
-                and x > self.__posxy[0]
-                and oldline[px:x] == newline[px + 1 : x + 1]
+                and x_coord > self.__posxy[0]
+                and oldline[px_pos:x_pos] == newline[px_pos + 1 : x_pos + 1]
             ):
-                x = px
-            self.__move(x, y)
+                x_pos = px_pos
+                x_coord = px_coord
+            self.__move(x_coord, y)
             self.__write_code(self.ich1)
-            self.__write(newline[x])
-            self.__posxy = x + 1, y
-        elif x < minlen and oldline[x + 1 :] == newline[x + 1 :]:
-            self.__move(x, y)
-            self.__write(newline[x])
-            self.__posxy = x + 1, y
+            self.__write(newline[x_pos])
+            self.__posxy = x_coord + character_width, y
+
+        # if it's a single character change in the middle of the line
+        elif x_coord < minlen and oldline[x_pos + 1 :] == newline[x_pos + 1 :] and wlen(oldline[x_pos]) == wlen(newline[x_pos]):
+            self.__move(x_coord, y)
+            self.__write(newline[x_pos])
+            self.__posxy = x_coord + character_width, y
+
+        # if this is the last character to fit in the line and we edit in the middle of the line
         elif (
             self.dch1
             and self.ich1
-            and len(newline) == self.width
-            and x < len(newline) - 2
-            and newline[x + 1 : -1] == oldline[x:-2]
+            and wlen(newline) == self.width
+            and x_coord < wlen(newline) - 2
+            and newline[x_pos + 1 : -1] == oldline[x_pos:-2]
         ):
             self.__hide_cursor()
             self.__move(self.width - 2, y)
             self.__posxy = self.width - 2, y
             self.__write_code(self.dch1)
-            self.__move(x, y)
+
+            self.__move(x_coord, y)
             self.__write_code(self.ich1)
-            self.__write(newline[x])
-            self.__posxy = x + 1, y
+            self.__write(newline[x_pos])
+            self.__posxy = x_coord + 1, y
+
         else:
             self.__hide_cursor()
-            self.__move(x, y)
-            if len(oldline) > len(newline):
+            self.__move(x_coord, y)
+            if wlen(oldline) > wlen(newline):
                 self.__write_code(self._el)
-            self.__write(newline[x:])
-            self.__posxy = len(newline), y
+            self.__write(newline[x_pos:])
+            self.__posxy = wlen(newline), y
 
         if "\x1b" in newline:
             # ANSI escape characters are present, so we can't assume