Cache parsed sources

Author: koxudaxi

src/datamodel_code_generator/__init__.py

@@ -9,11 +9,13 @@
 import contextlib
 import os
 import sys
-from collections import defaultdict
+from collections import OrderedDict, defaultdict
 from collections.abc import Callable, Iterator, Mapping
 from datetime import datetime, timezone
 from functools import lru_cache as _lru_cache
+from hashlib import sha256
 from pathlib import Path
+from threading import RLock
 from typing import (
     IO,
     TYPE_CHECKING,
@@ -94,6 +96,12 @@
 
 DEFAULT_BASE_CLASS: str = "pydantic.BaseModel"
 _IGNORED_TEXT_PREFIX_CHARS: frozenset[str] = frozenset({"\ufeff", " ", "\t", "\r", "\n"})
+_PARSER_SOURCE_DATA_CACHE_MAX_SIZE = 128
+_ParserSourceDataCacheKey: TypeAlias = tuple[Path, str, str, str]
+_ParserSourceDataSeenKey: TypeAlias = tuple[Path, str, str]
+_parser_source_data_cache: OrderedDict[_ParserSourceDataCacheKey, YamlValue] = OrderedDict()
+_parser_source_data_seen_keys: OrderedDict[_ParserSourceDataSeenKey, None] = OrderedDict()
+_parser_source_data_cache_lock = RLock()
 
 
 def load_yaml(stream: str | TextIO) -> YamlValue:
@@ -202,16 +210,76 @@ def load_data_from_path(path: Path, encoding: str) -> dict[str, YamlValue]:
 
     For file input: tries json.load() for .json files (more efficient than
     read_text + json.loads), falls back to YAML if JSON parsing fails
-    (e.g., trailing commas) or if content is not a dict. Uses YAML for all other extensions.
+    (e.g., trailing commas), and requires the parsed content to be a dict.
+    Uses YAML for all other extensions.
     """
+    result = _load_parser_source_data_from_path(path, encoding, cache_on_first_load=False)
+    if not isinstance(result, dict):
+        msg = f"Expected dict, got {type(result).__name__}"
+        raise TypeError(msg)
+    return result
+
+
+def _parser_source_data_yaml_backend(resolved_path: Path) -> str:
+    if resolved_path.suffix.lower() != ".json":
+        from datamodel_code_generator.util import get_yaml_backend  # noqa: PLC0415
+
+        return get_yaml_backend()
+    return ""
+
+
+def _load_parser_source_data_from_path(
+    path: Path,
+    encoding: str,
+    *,
+    cache_on_first_load: bool = True,
+) -> YamlValue:
+    resolved_path = path.resolve()
+    yaml_backend = _parser_source_data_yaml_backend(resolved_path)
+    seen_key = (resolved_path, encoding, yaml_backend)
+    with _parser_source_data_cache_lock:
+        use_cache = cache_on_first_load or seen_key in _parser_source_data_seen_keys
+        if not cache_on_first_load:
+            _parser_source_data_seen_keys[seen_key] = None
+            _parser_source_data_seen_keys.move_to_end(seen_key)
+            while len(_parser_source_data_seen_keys) > _PARSER_SOURCE_DATA_CACHE_MAX_SIZE:
+                _parser_source_data_seen_keys.popitem(last=False)
+
+    data = resolved_path.read_bytes()
+    if not use_cache:
+        return _load_parser_source_data_from_bytes(resolved_path, data, encoding)
+
+    digest = sha256(data).hexdigest()
+    cache_key = (resolved_path, digest, encoding, yaml_backend)
+    with _parser_source_data_cache_lock:
+        if cache_key in _parser_source_data_cache:
+            _parser_source_data_cache.move_to_end(cache_key)
+            return _parser_source_data_cache[cache_key]
+
+    parsed_data = _load_parser_source_data_from_bytes(resolved_path, data, encoding)
+    with _parser_source_data_cache_lock:
+        _parser_source_data_cache[cache_key] = parsed_data
+        _parser_source_data_cache.move_to_end(cache_key)
+        while len(_parser_source_data_cache) > _PARSER_SOURCE_DATA_CACHE_MAX_SIZE:
+            _parser_source_data_cache.popitem(last=False)
+    return parsed_data
+
+
+def _load_parser_source_data_from_bytes(path: Path, data: bytes, encoding: str) -> YamlValue:
     import json  # noqa: PLC0415
 
+    text = data.decode(encoding)
     if path.suffix.lower() == ".json":
-        with contextlib.suppress(json.JSONDecodeError), path.open(encoding=encoding) as f:
-            result = json.load(f)
-            if isinstance(result, dict):
-                return result
-    return load_yaml_dict_from_path(path, encoding)
+        with contextlib.suppress(json.JSONDecodeError):
+            return json.loads(text)
+
+    return load_yaml(text)
+
+
+def _clear_parser_source_data_cache() -> None:
+    with _parser_source_data_cache_lock:
+        _parser_source_data_cache.clear()
+        _parser_source_data_seen_keys.clear()
 
 
 @_lru_cache(maxsize=256)

src/datamodel_code_generator/_types/generate_config_dict.py

@@ -184,6 +184,7 @@ class GenerateConfigDict(TypedDict, closed=True):
     schema_version: NotRequired[str | None]
     schema_version_mode: NotRequired[VersionMode | None]
     external_ref_mapping: NotRequired[dict[str, str] | None]
+    enable_parsed_source_cache: NotRequired[bool]
 
 
 class ValidatorDefinition(TypedDict):

src/datamodel_code_generator/_types/parser_config_dicts.py

@@ -175,6 +175,7 @@ class ParserConfigDict(TypedDict):
     target_pydantic_version: NotRequired[TargetPydanticVersion | None]
     default_value_overrides: NotRequired[Mapping[str, Any] | None]
     external_ref_mapping: NotRequired[dict[str, str] | None]
+    enable_parsed_source_cache: NotRequired[bool]
 
 
 class GraphQLParserConfigDict(ParserConfigDict, closed=True):

src/datamodel_code_generator/config.py

@@ -212,6 +212,7 @@ class GenerateConfig(BaseModel):
     schema_version: str | None = None
     schema_version_mode: VersionMode | None = None
     external_ref_mapping: dict[str, str] | None = None
+    enable_parsed_source_cache: bool = False
 
 
 def _rebuild_generate_config() -> None:
@@ -351,6 +352,7 @@ class ParserConfig(BaseModel):
     target_pydantic_version: TargetPydanticVersion | None = None
     default_value_overrides: Mapping[str, Any] | None = None
     external_ref_mapping: dict[str, str] | None = None
+    enable_parsed_source_cache: bool = False
 
 
 class GraphQLParserConfig(ParserConfig):

src/datamodel_code_generator/parser/avro.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING, Any, NamedTuple, cast
+from typing import TYPE_CHECKING, Any, ClassVar, NamedTuple, cast
 
 from typing_extensions import Unpack
 
@@ -500,6 +500,7 @@ class AvroParser(JsonSchemaParser):
     """
 
     _config_class_name = "AvroParserConfig"
+    _cache_parsed_sources_from_path: ClassVar[bool] = False
 
     def __init__(
         self,

src/datamodel_code_generator/parser/base.py

@@ -49,6 +49,7 @@
     ReuseScope,
     YamlValue,
     _internal_utils,
+    _load_parser_source_data_from_path,
 )
 from datamodel_code_generator.enums import StrictTypes
 from datamodel_code_generator.format import (
@@ -949,11 +950,25 @@ class Source(BaseModel):
 
     path: Path
     text: str = ""
-    raw_data: dict[str, YamlValue] | None = None
+    raw_data: YamlValue | None = None
 
     @classmethod
-    def from_path(cls, path: Path, base_path: Path, encoding: str) -> Source:
+    def from_path(
+        cls,
+        path: Path,
+        base_path: Path,
+        encoding: str,
+        *,
+        enable_parsed_source_cache: bool = False,
+        keep_text: bool = False,
+    ) -> Source:
         """Create a Source from a file path relative to base_path."""
+        if enable_parsed_source_cache:
+            return cls(
+                path=path.relative_to(base_path),
+                text=path.read_text(encoding=encoding) if keep_text else "",
+                raw_data=_load_parser_source_data_from_path(path, encoding),
+            )
         return cls(
             path=path.relative_to(base_path),
             text=path.read_text(encoding=encoding),
@@ -1124,6 +1139,7 @@ def schema_features(self) -> SchemaFeaturesT:
 
     _config_class_name: ClassVar[str] = "ParserConfig"
     _cache_local_sources_during_parse: ClassVar[bool] = False
+    _cache_parsed_sources_from_path: ClassVar[bool] = False
 
     @classmethod
     def _get_config_class(cls) -> type[ParserConfig]:
@@ -1300,6 +1316,9 @@ def __init__(  # noqa: PLR0912, PLR0915
         self.source: str | Path | list[Path] | ParseResult | dict[str, YamlValue] = source
         self._cache_local_sources = False
         self._local_source_cache: tuple[Source, ...] | None = None
+        self.enable_parsed_source_cache: bool = (
+            config.enable_parsed_source_cache and self._cache_parsed_sources_from_path
+        )
         self.custom_template_dir = config.custom_template_dir
         self.extra_template_data: defaultdict[str, Any] = config.extra_template_data or defaultdict(dict)
         self.validators = config.validators
@@ -1519,12 +1538,30 @@ def _iter_source_uncached(self) -> Iterator[Source]:
                 if path.is_dir():
                     for p in sorted(path.rglob("*"), key=lambda p: p.name):
                         if p.is_file():
-                            yield Source.from_path(p, self.base_path, self.encoding)
+                            yield Source.from_path(
+                                p,
+                                self.base_path,
+                                self.encoding,
+                                enable_parsed_source_cache=self.enable_parsed_source_cache,
+                                keep_text=self.validation,
+                            )
                 else:
-                    yield Source.from_path(path, self.base_path, self.encoding)
+                    yield Source.from_path(
+                        path,
+                        self.base_path,
+                        self.encoding,
+                        enable_parsed_source_cache=self.enable_parsed_source_cache,
+                        keep_text=self.validation,
+                    )
             case list() as paths:  # pragma: no cover
                 for path in paths:
-                    yield Source.from_path(path, self.base_path, self.encoding)
+                    yield Source.from_path(
+                        path,
+                        self.base_path,
+                        self.encoding,
+                        enable_parsed_source_cache=self.enable_parsed_source_cache,
+                        keep_text=self.validation,
+                    )
             case _:
                 yield Source(
                     path=Path(self.source.path),

src/datamodel_code_generator/parser/jsonschema.py

@@ -639,6 +639,7 @@ class JsonSchemaParser(Parser["JSONSchemaParserConfig", "JsonSchemaFeatures"]):
     SCHEMA_PATHS: ClassVar[list[str]] = list(_DEFAULT_SCHEMA_PATHS)
     SCHEMA_OBJECT_TYPE: ClassVar[type[JsonSchemaObject]] = JsonSchemaObject
     _cache_local_sources_during_parse: ClassVar[bool] = True
+    _cache_parsed_sources_from_path: ClassVar[bool] = True
 
     COMPATIBLE_PYTHON_TYPES: ClassVar[dict[str, frozenset[str]]] = {
         "string": frozenset({"str", "String"}),
@@ -5352,7 +5353,12 @@ def _iter_local_source_paths(self) -> Iterator[Path]:
                 yield from ((self.base_path / path) for path in paths)
 
     def _load_source_dict(self, source: Source) -> dict[str, Any]:  # noqa: PLR6301
-        return dict(source.raw_data) if source.raw_data is not None else load_data(source.text)
+        if source.raw_data is None:
+            return load_data(source.text)
+        if not isinstance(source.raw_data, dict):
+            msg = f"Expected dict, got {type(source.raw_data).__name__}"
+            raise TypeError(msg)
+        return dict(source.raw_data)
 
     def _resolve_root_model_name(self, raw_obj: dict[str, Any]) -> tuple[str, bool]:
         title = raw_obj.get("title")
@@ -5397,17 +5403,11 @@ def parse_raw(self) -> None:
         """Parse all raw input sources into data models."""
         try:
             for source, path_parts in self._get_context_source_path_parts():
-                if source.raw_data is not None:
-                    raw_obj = source.raw_data
-                    if not isinstance(raw_obj, dict):  # pragma: no cover
-                        warn(f"{source.path} is empty or not a dict. Skipping this file", stacklevel=2)
-                        continue
-                else:
-                    try:
-                        raw_obj = load_data(source.text)
-                    except TypeError:
-                        warn(f"{source.path} is empty or not a dict. Skipping this file", stacklevel=2)
-                        continue
+                try:
+                    raw_obj = self._load_source_dict(source)
+                except TypeError:
+                    warn(f"{source.path} is empty or not a dict. Skipping this file", stacklevel=2)
+                    continue
                 self.raw_obj = raw_obj
                 obj_name, preserve_root_class_name = self._resolve_root_model_name(self.raw_obj)
                 self._parse_file(self.raw_obj, obj_name, path_parts, preserve_root_class_name=preserve_root_class_name)

src/datamodel_code_generator/parser/openapi.py

@@ -957,7 +957,7 @@ def parse_raw(self) -> None:
                 if self.validation:
                     warn_deprecated("cli.validation", stacklevel=2)
 
-                    if source.raw_data is not None:
+                    if source.raw_data is not None and not source.text:
                         warn(
                             "Warning: Validation was skipped for dict input. "
                             "The --validation option only works with file or text input.\n",

src/datamodel_code_generator/parser/protobuf.py

@@ -661,6 +661,7 @@ class ProtobufParser(JsonSchemaParser):
 
     _config_class_name: ClassVar[str] = "ProtobufParserConfig"
     _cache_local_sources_during_parse: ClassVar[bool] = False
+    _cache_parsed_sources_from_path: ClassVar[bool] = False
 
     def __init__(
         self,