Add generation dependency index (#3154)

* Add generation dependency index

* Fix generation store checks

* Fix generation store usage hook

* Improve generation coverage tests

Author: koxudaxi

.pre-commit-config.yaml

@@ -53,3 +53,9 @@ repos:
     language: system
     files: ^(src/datamodel_code_generator/parser/schema_version\.py|scripts/build_schema_docs\.py|docs/supported_formats\.md)$
     pass_filenames: false
+  - id: generation-store-usage
+    name: Check parser GenerationStore usage
+    entry: python scripts/check_generation_store_usage.py
+    language: python
+    files: ^src/datamodel_code_generator/parser/.*\.py$
+    pass_filenames: false

docs/development-contributing.md

@@ -84,6 +84,37 @@ assertions are also rejected. Only use `@pytest.mark.allow_direct_assert` for
 cases that cannot reasonably be expressed with the shared helpers, such as
 mock setup internals or narrow intermediate-state checks.
 
+## Parser generation state
+
+Parser code must keep generated output, naming order, parse order, and
+canonical model selection compatible with existing releases. `ModelResolver`
+remains the only authority for names. Do not move name generation or parser
+dispatch into the generation index.
+
+Use `GenerationStore` for parser-side mutations that affect generation facts:
+
+- Register parsed models with `generation_store.register_model(model)`.
+- Replace references with `generation_store.replace_data_type_ref(...)` or
+  `generation_store.detach_data_type_ref(...)`.
+- Replace field or nested data types with `replace_field_type(...)`,
+  `replace_nested_data_type(...)`, or `set_nested_data_types(...)`.
+- Replace fields or base classes with `set_fields(...)`,
+  `append_field(...)`, `insert_field(...)`, `remove_field(...)`, or
+  `set_base_classes(...)`.
+- Move or rename generated models with `move_model(...)`,
+  `rename_model(...)`, or `update_model_reference(...)`.
+
+Parser code should not directly assign `data_type.reference`, `model.fields`,
+`model.base_classes`, `model.reference.name`, or call reference child mutation
+helpers. Read dependency information through `GenerationIndex` queries instead
+of using `Reference.children` as a reverse index.
+
+The local pre-commit hook `generation-store-usage` runs
+`scripts/check_generation_store_usage.py` and rejects common parser mutations
+that bypass the store. If a new mutation helper is needed, add the method to
+`GENERATION_STORE_MUTATION_METHODS` and update the checker tests so contributors
+get an actionable message.
+
 ## ➕ Adding a New CLI Option
 
 When adding a new CLI option to `datamodel-code-generator`, follow these steps:

scripts/check_generation_store_usage.py

@@ -0,0 +1,244 @@
+"""Reject parser mutations that bypass GenerationStore."""
+
+from __future__ import annotations
+
+import argparse
+import ast
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TypeAlias
+
+ROOT = Path(__file__).resolve().parents[1]
+
+DEFAULT_TARGET = Path("src/datamodel_code_generator/parser")
+GENERATION_STORE_MODULE = Path("src/datamodel_code_generator/parser/generation.py")
+
+MUTATION_ATTRIBUTES = frozenset({"base_classes", "data_type", "fields", "reference"})
+MODEL_COLLECTION_ATTRIBUTES = frozenset({"models", "results"})
+MODEL_COLLECTION_HELPERS = {"append": "register_model"}
+REFERENCE_METHODS = frozenset(
+    {
+        "_invalidate_after_mutation",
+        "_refresh_after_mutation",
+        "remove_reference",
+        "replace_children_in_models",
+        "replace_children_references",
+        "replace_data_type",
+        "replace_reference",
+        "refresh_after_mutation",
+        "set_reference_path",
+    },
+)
+SEQUENCE_MUTATORS = frozenset({"append", "clear", "extend", "insert", "pop", "remove"})
+SequenceOwner: TypeAlias = str
+
+
+def _literal_string_set(node: ast.AST) -> frozenset[str]:
+    if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "frozenset" and node.args:
+        node = node.args[0]
+    if not isinstance(node, ast.Set):
+        msg = "GENERATION_STORE_MUTATION_METHODS must be a literal frozenset of strings"
+        raise TypeError(msg)
+
+    values: set[str] = set()
+    for item in node.elts:
+        if not isinstance(item, ast.Constant) or not isinstance(item.value, str):
+            msg = "GENERATION_STORE_MUTATION_METHODS must contain only literal strings"
+            raise TypeError(msg)
+        values.add(item.value)
+    return frozenset(values)
+
+
+def _load_generation_store_mutation_methods() -> frozenset[str]:
+    generation_store_module = ROOT / GENERATION_STORE_MODULE
+    tree = ast.parse(generation_store_module.read_text(), filename=str(generation_store_module))
+    for node in tree.body:
+        if (
+            isinstance(node, ast.AnnAssign)
+            and isinstance(node.target, ast.Name)
+            and node.target.id == "GENERATION_STORE_MUTATION_METHODS"
+            and node.value is not None
+        ):
+            return _literal_string_set(node.value)
+        if isinstance(node, ast.Assign):
+            for target in node.targets:
+                if isinstance(target, ast.Name) and target.id == "GENERATION_STORE_MUTATION_METHODS":
+                    return _literal_string_set(node.value)
+    msg = "GENERATION_STORE_MUTATION_METHODS was not found"
+    raise RuntimeError(msg)
+
+
+GENERATION_STORE_MUTATION_METHODS = _load_generation_store_mutation_methods()
+
+
+@dataclass(frozen=True)
+class Violation:
+    """A direct parser mutation that should go through GenerationStore."""
+
+    path: Path
+    line: int
+    column: int
+    message: str
+
+    def format(self) -> str:
+        """Format violation for CLI output."""
+        return f"{self.path}:{self.line}:{self.column}: {self.message}"
+
+
+def _attribute_chain(node: ast.AST) -> tuple[str, ...]:
+    names: list[str] = []
+    current = node
+    while isinstance(current, ast.Attribute):
+        names.append(current.attr)
+        current = current.value
+    if isinstance(current, ast.Name):
+        names.append(current.id)
+    names.reverse()
+    return tuple(names)
+
+
+def _sequence_owner(node: ast.AST) -> SequenceOwner | None:
+    if isinstance(node, ast.Attribute) and node.attr in {"base_classes", "fields", *MODEL_COLLECTION_ATTRIBUTES}:
+        return node.attr
+    if isinstance(node, ast.Subscript):
+        return _sequence_owner(node.value)
+    return None
+
+
+class GenerationStoreUsageVisitor(ast.NodeVisitor):
+    """AST visitor for parser code that must use GenerationStore mutation APIs."""
+
+    def __init__(self, path: Path) -> None:
+        """Create a visitor for ``path``."""
+        self.path = path
+        self.violations: list[Violation] = []
+
+    def visit_Assign(self, node: ast.Assign) -> None:
+        """Check assignment targets."""
+        for target in node.targets:
+            self._check_assignment_target(target)
+        self.generic_visit(node)
+
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        """Check annotated assignment targets."""
+        self._check_assignment_target(node.target)
+        self.generic_visit(node)
+
+    def visit_AugAssign(self, node: ast.AugAssign) -> None:
+        """Check augmented assignment targets."""
+        self._check_assignment_target(node.target)
+        self.generic_visit(node)
+
+    def visit_Call(self, node: ast.Call) -> None:
+        """Check direct mutating method calls."""
+        if isinstance(node.func, ast.Attribute):
+            method = node.func.attr
+            owner = _sequence_owner(node.func.value)
+            if owner and method in SEQUENCE_MUTATORS:
+                self._add(
+                    node,
+                    f"use GenerationStore.{self._sequence_helper(owner, method)}() instead of mutating {owner}",
+                )
+            elif method in REFERENCE_METHODS:
+                self._add(node, f"use GenerationStore API instead of {method}()")
+        self.generic_visit(node)
+
+    def visit_Attribute(self, node: ast.Attribute) -> None:
+        """Check direct ``Reference.children`` reads."""
+        if node.attr == "children" and isinstance(node.ctx, ast.Load):
+            self._add(node, "use GenerationIndex reverse-edge queries instead of Reference.children")
+        self.generic_visit(node)
+
+    def _check_assignment_target(self, target: ast.AST) -> None:
+        owner = _sequence_owner(target)
+        if owner:
+            helper = self._sequence_helper(owner, "assign")
+            self._add(target, f"use GenerationStore.{helper}() instead of assigning {owner}")
+            return
+
+        if not isinstance(target, ast.Attribute):
+            return
+
+        chain = _attribute_chain(target)
+        if not chain:
+            return
+
+        if target.attr in MUTATION_ATTRIBUTES or (
+            isinstance(target.value, ast.Attribute) and target.value.attr == "reference"
+        ):
+            self._add(target, f"use GenerationStore API instead of assigning {'.'.join(chain[-2:])}")
+
+    def _add(self, node: ast.AST, message: str) -> None:
+        self.violations.append(
+            Violation(
+                path=self.path,
+                line=getattr(node, "lineno", 0),
+                column=getattr(node, "col_offset", 0) + 1,
+                message=message,
+            )
+        )
+
+    @staticmethod
+    def _sequence_helper(owner: SequenceOwner, method: str) -> str:
+        if owner in MODEL_COLLECTION_ATTRIBUTES:
+            helper = MODEL_COLLECTION_HELPERS.get(method, "register_model")
+            return _ensure_generation_store_api(helper)
+        if owner == "fields":
+            helper = {
+                "assign": "set_fields",
+                "append": "append_field",
+                "insert": "insert_field",
+                "remove": "remove_field",
+            }.get(method, "set_fields")
+            return _ensure_generation_store_api(helper)
+        return _ensure_generation_store_api("set_base_classes")
+
+
+def _ensure_generation_store_api(method_name: str) -> str:
+    if method_name not in GENERATION_STORE_MUTATION_METHODS:
+        msg = f"{method_name!r} is not a registered GenerationStore mutation API"
+        raise RuntimeError(msg)
+    return method_name
+
+
+def iter_python_files(paths: list[Path]) -> list[Path]:
+    """Return checked Python files below ``paths``."""
+    files: list[Path] = []
+    for path in paths:
+        if path.is_dir():
+            files.extend(sorted(path.rglob("*.py")))
+        elif path.suffix == ".py":
+            files.append(path)
+    generation_store_module = GENERATION_STORE_MODULE.resolve()
+    return [path for path in files if path.resolve() != generation_store_module]
+
+
+def check_paths(paths: list[Path]) -> list[Violation]:
+    """Check paths and return all GenerationStore usage violations."""
+    violations: list[Violation] = []
+    for path in iter_python_files(paths):
+        tree = ast.parse(path.read_text(), filename=str(path))
+        visitor = GenerationStoreUsageVisitor(path)
+        visitor.visit(tree)
+        violations.extend(visitor.violations)
+    return violations
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entry point."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("paths", nargs="*", type=Path, default=[DEFAULT_TARGET])
+    args = parser.parse_args(argv)
+
+    violations = check_paths(args.paths)
+    if violations:
+        print("GenerationStore usage violations found:", file=sys.stderr)
+        for violation in violations:
+            print(violation.format(), file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())