diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index e39193f51de..cafec8d249b 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -122,6 +122,7 @@ import sys import tokenize from collections.abc import Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass +from itertools import groupby from multiprocessing import Pool from pathlib import Path from types import MappingProxyType @@ -880,18 +881,35 @@ def _bool_constant(value: ast.expr) -> bool | None: return value.value if isinstance(value, ast.Constant) and isinstance(value.value, bool) else None -def _module_assignment_items(tree: ast.AST) -> Iterator[tuple[str, ast.expr]]: +def _module_assignment_items(tree: ast.AST) -> Iterator[tuple[str, tuple[int, ast.expr]]]: if not isinstance(tree, ast.Module): return for stmt in tree.body: if isinstance(stmt, ast.Assign): - yield from ((target.id, stmt.value) for target in stmt.targets if isinstance(target, ast.Name)) + yield from ( + (target.id, (stmt.lineno, stmt.value)) for target in stmt.targets if isinstance(target, ast.Name) + ) elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name) and stmt.value is not None: - yield stmt.target.id, stmt.value + yield stmt.target.id, (stmt.lineno, stmt.value) -def _model_config_frozen(value: ast.expr, module_assignments: Mapping[str, ast.expr]) -> bool | None: - config: Final = module_assignments.get(value.id) if isinstance(value, ast.Name) else value +def _model_config_frozen( + value: ast.expr, + module_assignments: Mapping[str, tuple[tuple[int, ast.expr], ...]], + class_lineno: int, +) -> bool | None: + config: Final = ( + next( + ( + expression + for lineno, expression in reversed(module_assignments.get(value.id, ())) + if lineno < class_lineno + ), + None, + ) + if isinstance(value, ast.Name) + else value + ) if isinstance(config, ast.Call) and _head_name(config.func) in PYDANTIC_CONFIG_FACTORIES: flags: Final = tuple(_bool_constant(kw.value) for kw in config.keywords if kw.arg == "frozen") return flags[-1] if flags else None @@ -921,28 +939,44 @@ def _config_class_frozen(node: ast.ClassDef) -> bool | None: return flags[-1] if flags else None -def _stmt_frozen_flag(stmt: ast.stmt, module_assignments: Mapping[str, ast.expr]) -> bool | None: +def _stmt_frozen_flag( + stmt: ast.stmt, + module_assignments: Mapping[str, tuple[tuple[int, ast.expr], ...]], + class_lineno: int, +) -> bool | None: config_value: Final = _assigns_name(stmt, "model_config") if config_value is not None: - return _model_config_frozen(config_value, module_assignments) + return _model_config_frozen(config_value, module_assignments, class_lineno) if isinstance(stmt, ast.ClassDef) and stmt.name == "Config": return _config_class_frozen(stmt) return None -def _class_frozen_override(cls: ast.ClassDef, module_assignments: Mapping[str, ast.expr]) -> bool | None: +def _class_frozen_override( + cls: ast.ClassDef, + module_assignments: Mapping[str, tuple[tuple[int, ast.expr], ...]], + class_lineno: int, +) -> bool | None: keyword_flags: Final = tuple(_bool_constant(keyword.value) for keyword in cls.keywords if keyword.arg == "frozen") if keyword_flags: return keyword_flags[-1] - body_flags: Final = tuple(_stmt_frozen_flag(stmt, module_assignments) for stmt in cls.body) + body_flags: Final = tuple(_stmt_frozen_flag(stmt, module_assignments, class_lineno) for stmt in cls.body) return next((flag for flag in reversed(body_flags) if flag is not None), None) def iter_pydantic_violations(path: Path, tree: ast.AST) -> Iterator[Violation]: - module_assignments: Final = MappingProxyType(dict(_module_assignment_items(tree))) + module_assignments: Final = MappingProxyType( + { + name: tuple(binding for _, binding in assignments) + for name, assignments in groupby( + sorted(_module_assignment_items(tree), key=lambda item: item[0]), + key=lambda item: item[0], + ) + } + ) models: Final = _pydantic_classes(tree) bases_of: Final = {cls: _base_names(cls) for cls in models} - override_of: Final = {cls: _class_frozen_override(cls, module_assignments) for cls in models} + override_of: Final = {cls: _class_frozen_override(cls, module_assignments, cls.lineno) for cls in models} def frozen(known: frozenset[str]) -> frozenset[str]: grown: Final = known | frozenset( diff --git a/tests/unit/test_check_type_discipline.py b/tests/unit/test_check_type_discipline.py index 60852654776..abb7ea1408c 100644 --- a/tests/unit/test_check_type_discipline.py +++ b/tests/unit/test_check_type_discipline.py @@ -611,14 +611,37 @@ def test_shared_frozen_configdict_constant_is_clean(tmp_path): assert "LIT015" not in _codes(tmp_path, src) -def test_last_shared_non_frozen_configdict_assignment_is_flagged(tmp_path): +def test_mutable_config_before_class_is_not_hidden_by_frozen_reassignment(tmp_path): src: Final = ( "from pydantic import BaseModel, ConfigDict\n" - "_RESPONSE_CONFIG = ConfigDict(frozen=True)\n" "_RESPONSE_CONFIG = ConfigDict(frozen=False)\n" "class Foo(BaseModel):\n" " model_config = _RESPONSE_CONFIG\n" " x: int\n" + "_RESPONSE_CONFIG = ConfigDict(frozen=True)\n" + ) + assert "LIT015" in _codes(tmp_path, src) + + +def test_frozen_config_before_class_survives_mutable_reassignment(tmp_path): + src: Final = ( + "from pydantic import BaseModel, ConfigDict\n" + "_RESPONSE_CONFIG = ConfigDict(frozen=True)\n" + "class Foo(BaseModel):\n" + " model_config = _RESPONSE_CONFIG\n" + " x: int\n" + "_RESPONSE_CONFIG = ConfigDict(frozen=False)\n" + ) + assert "LIT015" not in _codes(tmp_path, src) + + +def test_config_assigned_only_after_class_is_unresolved(tmp_path): + src: Final = ( + "from pydantic import BaseModel, ConfigDict\n" + "class Foo(BaseModel):\n" + " model_config = _RESPONSE_CONFIG\n" + " x: int\n" + "_RESPONSE_CONFIG = ConfigDict(frozen=True)\n" ) assert "LIT015" in _codes(tmp_path, src)