Cache builtin formatter source lines (#3552)

* Cache builtin formatter source lines

* Satisfy source segment typing

* Avoid source lines cache global assignment

* Cover builtin formatter location fallbacks

Author: koxudaxi

src/datamodel_code_generator/_builtin_formatter.py

@@ -9,7 +9,7 @@
 from collections import defaultdict
 from enum import IntEnum
 from io import StringIO
-from typing import TYPE_CHECKING, Any, TypeGuard
+from typing import TYPE_CHECKING, Any, TypeGuard, cast
 
 from datamodel_code_generator.util import load_toml
 
@@ -23,6 +23,12 @@ class PythonVersion(Protocol):
         def version_key(self) -> tuple[int, int]:
             raise NotImplementedError
 
+    class _SourceLocationNode(Protocol):
+        lineno: int
+        end_lineno: int | None
+        col_offset: int
+        end_col_offset: int | None
+
 
 class _AliasSortCategory(IntEnum):
     CONSTANT = 0
@@ -47,6 +53,7 @@ class _ImportCategory(IntEnum):
 STRING_PREFIX_PATTERN = re.compile(r"(?i)^([rubf]*)(\"\"\"|'''|\"|')")
 PEP695_TYPE_ALIAS_START_PATTERN = re.compile(r"^(?P<indent>\s*)type\s+(?P<target>[A-Za-z_]\w*(?:\[.*?\])?)\s*=")
 PEP695_TYPE_ALIAS_PLACEHOLDER = "__datamodel_codegen_builtin_type_alias__"
+_SOURCE_LINES_CACHE: list[tuple[str, list[str]]] = []
 
 
 def _is_valid_builtin_line_length(line_length: Any) -> TypeGuard[int]:
@@ -417,8 +424,70 @@ def _is_type_checking_if(node: ast.AST) -> TypeGuard[ast.If]:
     return isinstance(node, ast.If) and _is_name_or_attr(node.test, "TYPE_CHECKING")
 
 
+def _splitlines_no_ff(source: str) -> list[str]:
+    """Split source lines like the Python parser, without form-feed splitting."""
+    index = 0
+    lines: list[str] = []
+    next_line = ""
+    while index < len(source):
+        character = source[index]
+        next_line += character
+        index += 1
+
+        if character == "\r" and index < len(source) and source[index] == "\n":
+            next_line += "\n"
+            index += 1
+        if character in "\r\n":
+            lines.append(next_line)
+            next_line = ""
+
+    if next_line:
+        lines.append(next_line)
+    return lines
+
+
+def _source_lines(source: str) -> list[str]:
+    if _SOURCE_LINES_CACHE:
+        cached_source, cached_lines = _SOURCE_LINES_CACHE[0]
+        if source is cached_source or source == cached_source:
+            return cached_lines
+
+    lines = _splitlines_no_ff(source)
+    _SOURCE_LINES_CACHE[:] = [(source, lines)]
+    return lines
+
+
+def _source_segment_from_cached_lines(source: str, node: ast.AST) -> str | None:
+    if not (
+        hasattr(node, "lineno")
+        and hasattr(node, "end_lineno")
+        and hasattr(node, "col_offset")
+        and hasattr(node, "end_col_offset")
+    ):
+        return None
+
+    location_node = cast("_SourceLocationNode", node)
+    lineno = location_node.lineno
+    end_lineno = location_node.end_lineno
+    col_offset = location_node.col_offset
+    end_col_offset = location_node.end_col_offset
+    if end_lineno is None or end_col_offset is None:
+        return None
+
+    lineno -= 1
+    end_lineno -= 1
+
+    lines = _source_lines(source)
+    if end_lineno == lineno:
+        return lines[lineno].encode()[col_offset:end_col_offset].decode()
+
+    first = lines[lineno].encode()[col_offset:].decode()
+    last = lines[end_lineno].encode()[:end_col_offset].decode()
+    return "".join([first, *lines[lineno + 1 : end_lineno], last])
+
+
 def _source_segment(source: str, node: ast.AST) -> str:
-    return ast.get_source_segment(source, node) or ast.unparse(node)
+    return _source_segment_from_cached_lines(source, node) or ast.unparse(node)
 
 
 def _inline_source_segment(source: str, node: ast.AST) -> str:

tests/test_format.py

@@ -1052,6 +1052,71 @@ def test_builtin_formatter_pep695_type_alias_replacement_edges() -> None:
     ]
 
 
+def test_builtin_formatter_source_segment_matches_ast() -> None:
+    """Source extraction should match ast.get_source_segment, including UTF-8 byte offsets."""
+    source = (
+        "from typing import Annotated\r\n"
+        "from pydantic import Field\n"
+        "\n"
+        "EXTRA = {'emoji': '🍣'}\n"
+        "\n"
+        "class Model:\n"
+        '    """説明."""\n'
+        "    field: Annotated[\n"
+        "        str,\n"
+        "        Field(default='é', description='長い説明'),\n"
+        "    ]\n"
+        "    data = {\n"
+        "        'emoji': '🍣',\n"
+        "        **EXTRA,\n"
+        "    }\n"
+        "\n"
+        "lambda_value = lambda: '空'\n"
+    )
+
+    for node in ast.walk(ast.parse(source)):
+        expected = ast.get_source_segment(source, node)
+        if not expected:
+            continue
+
+        assert builtin_formatter._source_segment(source, node) == expected
+
+
+def test_builtin_formatter_source_segment_reuses_split_source_lines(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Repeated source extraction should split one source object once."""
+    source = "class Model:\n    field: str = 'é'\n    other: str = '🍣'\n"
+    nodes = [node for node in ast.walk(ast.parse(source)) if ast.get_source_segment(source, node)]
+    split_call_count = 0
+    original_splitlines = builtin_formatter._splitlines_no_ff
+
+    def spy_splitlines(value: str) -> list[str]:
+        nonlocal split_call_count
+        split_call_count += 1
+        return original_splitlines(value)
+
+    monkeypatch.setattr(builtin_formatter, "_SOURCE_LINES_CACHE", [])
+    monkeypatch.setattr(builtin_formatter, "_splitlines_no_ff", spy_splitlines)
+
+    for node in nodes:
+        builtin_formatter._source_segment(source, node)
+
+    assert split_call_count == 1
+
+
+def test_builtin_formatter_source_segment_handles_missing_locations() -> None:
+    """Source extraction should fall back when a node lacks complete location data."""
+    source = "field: str\n"
+    node_without_location = ast.Load()
+    node_without_end_location = ast.Name(id="field", ctx=ast.Load())
+    node_without_end_location.lineno = 1
+    node_without_end_location.col_offset = 0
+    node_without_end_location.end_lineno = None
+    node_without_end_location.end_col_offset = None
+
+    assert builtin_formatter._source_segment_from_cached_lines(source, node_without_location) is None
+    assert builtin_formatter._source_segment_from_cached_lines(source, node_without_end_location) is None
+
+
 def test_apply_builtin_formatter_parenthesizes_short_annotated_default() -> None:
     """Test built-in formatter matches black for short overflowing Annotated defaults."""
     code = (