Merge pull request #872 from python-rope/lieryan-implement-ast-lines-adapter

Fix unicode handling in patched ast

Author: lieryan

CHANGELOG.md

@@ -2,11 +2,12 @@
 
 - #850 Update and pin black version in pre-commit and Github Actions
 - #851 Bump supported python version to up to Python 3.14
-- #852 Implement patchedast handlers for TypeAlias 
+- #852 Implement patchedast handlers for TypeAlias
 - #853 Implement patchedast handlers TypeVar
 - #847 Avoid printing autoimport syntax errors (@yangfan-yf-yf)
 - #623, #819, #863 Support MatchOr, MatchSequence, MatchStar (@jheld, @lieryan)
 - #870 Add default implementation for is_dir() (@lieryan)
+- #872 Fix unicode handling in patchedast (@lieryan)
 
 # Release 1.14.0
 

rope/base/codeanalyze.py

@@ -70,24 +70,147 @@ def get_line_start(self, lineno):
     def get_line_end(self, lineno):
         return self.starts[lineno] - 1
 
-    def __getitem__(self, subscript):
-        start_offset = self._calculate_offset(subscript.start)
-        stop_offset = self._calculate_offset(subscript.stop)
-        return self.code[start_offset:stop_offset]
-
-    def _calculate_offset(self, coord: tuple[int, int]) -> int:
-        lineno, col_offset = coord
-        lineno = self._clamp(0, self.length(), lineno)
-        col_offset = self._clamp(
-            0,
-            self.get_line_end(lineno) - self.get_line_start(lineno),
-            col_offset,
+
+class ASTLinesAdapter:
+    """
+    Convert between stdlib's ast `col_offset` and rope's `region` offsets.
+
+    ast `col_offset` are given in byte offset of the utf-8 source code from
+    the CPython interpreter, this is not necessarily identical to the str
+    offset (i.e. unicode codepoint offset) that rope uses.
+
+    Semantics are intended to match ``ast.get_source_segment`` exactly,
+    including its quirks:
+
+    - Line numbers are 1-indexed (``node.lineno`` / ``node.end_lineno``),
+      matching the ``ast`` module.
+    - Column offsets are UTF-8 *byte* offsets into the line, not character
+      offsets -- this matters for any line containing non-ASCII characters.
+    - Line splitting uses the same boundaries the CPython parser uses
+      (``\\n``, ``\\r\\n``, ``\\r``) rather than the broader set recognized by
+      ``str.splitlines()`` (form feed, vertical tab, etc.), since those extra
+      separators don't correspond to a new line as far as ``lineno`` is
+      concerned.
+    """
+
+    class LineNumberOutOfRange(IndexError):
+        pass
+
+    # Same pattern rule as `ast._line_pattern` used by `ast._splitlines_no_ff`
+    # to split lines in a way that mimics how the CPython parser splits source
+    # code.
+    LINE_PATTERN_STR = re.compile("(.*?(?:\r\n|\n|\r|$))")
+    LINE_PATTERN_BYTES = re.compile(b"(.*?(?:\r\n|\n|\r|$))")
+
+    def __init__(self, source: str):
+        self._code_str: str = source
+        self._code_bytes: bytes = source.encode("utf-8")
+        self._starts_str: list[int] = self._initialize_line_starts(
+            self.LINE_PATTERN_STR, self._code_str
         )
+        self._starts_bytes: list[int] = self._initialize_line_starts(
+            self.LINE_PATTERN_BYTES, self._code_bytes
+        )
+
+    @classmethod
+    def _initialize_line_starts(cls, pattern, source) -> list[int]:
+        matches = list(pattern.finditer(source))
+        if len(matches) > 1 and matches[-1].start() == matches[-1].end():
+            matches.pop()
+        starts = [m.start() for m in matches]
+        if source[-1:] in [b"\n", b"\r", "\n", "\r"]:
+            starts.append(len(source))
+        else:
+            starts.append(len(source) + 1)
+        return starts
+
+    def __len__(self):
+        return len(self._starts_bytes) - 1
+
+    def _validate_lineno(self, lineno: int) -> int:
+        if not (0 < lineno <= len(self)):
+            raise ASTLinesAdapter.LineNumberOutOfRange(
+                f"Line number out of range: {lineno}"
+            )
+        return lineno - 1
+
+    def _get_line_text(self, line_idx: int, col_offset: int | None = None) -> str:
+        return self._get_line_bytes(line_idx, col_offset).decode("utf-8")
+
+    def _get_line_bytes(self, line_idx: int, col_offset: int | None = None) -> bytes:
+        start = self._starts_bytes[line_idx]
+        end = self._starts_bytes[line_idx + 1]
+        if col_offset is not None:
+            end = min(end, start + col_offset)
+        return self._code_bytes[start:end].rstrip(b"\r\n")
+
+    def _is_line_ascii_only(self, line_idx: int):
+        start_str = self._starts_str[line_idx]
+        end_str = self._starts_str[line_idx + 1]
+        start_bytes = self._starts_bytes[line_idx]
+        end_bytes = self._starts_bytes[line_idx + 1]
+        # str length == bytes length iff they're ascii-only because any
+        # non-ascii characters would be at least 2 bytes
+        return (end_str - start_str) == (end_bytes - start_bytes)
+
+    def _line_region_offset(self, line_idx: int, col_offset: int) -> int:
+        """str offset relative to the start of line"""
+        if col_offset == 0 or self._is_line_ascii_only(line_idx):
+            return col_offset
+        prefix = self._get_line_text(line_idx, col_offset)
+        return len(prefix)
+
+    def _absolute_region_offset(self, line_idx: int, col_offset: int) -> int:
+        """str offset relative to the start of file"""
+        line_offset = self._line_region_offset(line_idx, col_offset)
+        return self._starts_str[line_idx] + line_offset
+
+    @classmethod
+    def _pad_whitespace(cls, source: str) -> str:
+        """Equivalent to `ast._pad_whitespace(source)`"""
+        return "".join(c if c in "\f\t" else " " for c in source)
+
+    def get_line(self, lineno: int) -> str:
+        line_idx = self._validate_lineno(lineno)
+        return self._get_line_text(line_idx)
+
+    def get_source_segment(self, node, *, padded: bool = False) -> str | None:
+        """Equivalent to ``ast.get_source_segment(source, node, padded=padded)``.
+
+        ``ast.get_source_segment(source, node)`` is convenient, but every call
+        re-splits the *entire* source string into lines (via the private
+        ``ast._splitlines_no_ff``) and re-encodes whatever lines it touches to
+        UTF-8. If you're pulling source segments for many nodes out of the same
+        file (e.g. walking a whole AST), that's a lot of repeated O(n) work.
+        """
+
+        region_start, region_end = self[node]
+        if region_start is None and region_end is None:
+            return None
+        source_segment = self._code_str[region_start:region_end]
+
+        if padded and node.lineno != node.end_lineno:
+            line_idx = node.lineno - 1
+            prefix = self._get_line_text(line_idx, node.col_offset)
+            padding = self._pad_whitespace(prefix)
+            source_segment = padding + source_segment
 
-        return self.get_line_start(lineno) + col_offset
+        return source_segment
 
-    def _clamp(self, min_value, max_value, value):
-        return max(min_value, min(max_value, value))
+    def __getitem__(self, node) -> tuple[int, int] | tuple[None, None]:
+        try:
+            if node.end_lineno is None or node.end_col_offset is None:
+                return (None, None)
+            line_idx = node.lineno - 1
+            end_line_idx = node.end_lineno - 1
+            col_offset = node.col_offset
+            end_col_offset = node.end_col_offset
+        except AttributeError:
+            return (None, None)
+
+        region_start = self._absolute_region_offset(line_idx, col_offset)
+        region_end = self._absolute_region_offset(end_line_idx, end_col_offset)
+        return (region_start, region_end)
 
 
 class ArrayLinesAdapter:

rope/refactor/patchedast.py

@@ -67,6 +67,7 @@ def __init__(self, source, children=False):
         self.source = _Source(source)
         self.children = children
         self.lines = codeanalyze.SourceLinesAdapter(source)
+        self.ast_adapter = codeanalyze.ASTLinesAdapter(source)
         self.children_stack = []
 
     Number = object()
@@ -217,7 +218,8 @@ def _find_next_statement_start(self):
         for children in reversed(self.children_stack):
             for child in children:
                 if isinstance(child, ast.stmt):
-                    return child.col_offset + self.lines.get_line_start(child.lineno)
+                    start, _ = self.ast_adapter[child]
+                    return start
         return len(self.source.source)
 
     def _join(self, iterable, separator):
@@ -576,11 +578,8 @@ def _If(self, node):
     def _is_elif(self, node):
         if not isinstance(node, ast.If):
             return False
-        offset = self.lines.get_line_start(node.lineno) + node.col_offset
-        word = self.source[offset : offset + 4]
-        # XXX: This is a bug; the offset does not point to the first
-        alt_word = self.source[offset - 5 : offset - 1]
-        return "elif" in (word, alt_word)
+        start, end = self.ast_adapter[node]
+        return "elif" in self.source[start : start + 4]
 
     def _IfExp(self, node):
         return self._handle(node, [node.body, "if", node.test, "else", node.orelse])
@@ -813,27 +812,23 @@ def _MatchSequence(self, node):
                 *closing_paren,
             ]
         else:
-            node_start = (node.lineno, node.col_offset)
-            node_end = (node.end_lineno, node.end_col_offset)
-            children = [self.lines[node_start:node_end]]
+            empty_tuple = self.ast_adapter.get_source_segment(node)
+            children = [empty_tuple]
         self._handle(node, children)
 
     def _get_surrounding_parens(self, node: ast.MatchSequence):
-        node_start = (node.lineno, node.col_offset)
-        first_pattern_start = (node.patterns[0].lineno, node.patterns[0].col_offset)
-        opening_paren = self.lines[node_start:first_pattern_start].strip()
+        node_start, node_end = self.ast_adapter[node]
+        first_pattern_start, _ = self.ast_adapter[node.patterns[0]]
+        _, last_pattern_end = self.ast_adapter[node.patterns[-1]]
+        opening_paren = self.source[node_start:first_pattern_start].strip()
+        closing_paren = self.source[last_pattern_end:node_end].strip()
+
         if opening_paren not in ["[", "(", ""]:
             warnings.warn(
                 f"Unexpected character in MatchSequence's opening_paren <{opening_paren}>; please report!",
                 RuntimeWarning,
             )
 
-        last_pattern_end = (
-            node.patterns[-1].end_lineno,
-            node.patterns[-1].end_col_offset,
-        )
-        node_end = (node.end_lineno, node.end_col_offset)
-        closing_paren = self.lines[last_pattern_end:node_end].strip()
         if closing_paren not in ["]", ")", ""]:
             warnings.warn(
                 f"Unexpected character in MatchSequence's closing_paren <{closing_paren}>; please report!",
@@ -1015,9 +1010,6 @@ def find_backwards(self, pattern, offset):
     def __getitem__(self, index):
         return self.source[index]
 
-    def __getslice__(self, i, j):
-        return self.source[i:j]
-
     def _get_number_pattern(self):
         # HACK: It is merely an approaximation and does the job
         integer = r"\-?(0[xo][\da-fA-F]+|\d+)"