diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index 02555f7d4e5..e39193f51de 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -880,14 +880,25 @@ def _bool_constant(value: ast.expr) -> bool | None: return value.value if isinstance(value, ast.Constant) and isinstance(value.value, bool) else None -def _model_config_frozen(value: ast.expr) -> bool | None: - if isinstance(value, ast.Call) and _head_name(value.func) in PYDANTIC_CONFIG_FACTORIES: - flags: Final = tuple(_bool_constant(kw.value) for kw in value.keywords if kw.arg == "frozen") +def _module_assignment_items(tree: ast.AST) -> Iterator[tuple[str, 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)) + elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name) and stmt.value is not None: + yield stmt.target.id, 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 + 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 - if isinstance(value, ast.Dict): + if isinstance(config, ast.Dict): flags: Final = tuple( _bool_constant(item) - for key, item in zip(value.keys, value.values) + for key, item in zip(config.keys, config.values) if isinstance(key, ast.Constant) and key.value == "frozen" ) return flags[-1] if flags else None @@ -910,27 +921,28 @@ def _config_class_frozen(node: ast.ClassDef) -> bool | None: return flags[-1] if flags else None -def _stmt_frozen_flag(stmt: ast.stmt) -> bool | None: +def _stmt_frozen_flag(stmt: ast.stmt, module_assignments: Mapping[str, ast.expr]) -> bool | None: config_value: Final = _assigns_name(stmt, "model_config") if config_value is not None: - return _model_config_frozen(config_value) + return _model_config_frozen(config_value, module_assignments) if isinstance(stmt, ast.ClassDef) and stmt.name == "Config": return _config_class_frozen(stmt) return None -def _class_frozen_override(cls: ast.ClassDef) -> bool | None: +def _class_frozen_override(cls: ast.ClassDef, module_assignments: Mapping[str, ast.expr]) -> 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(flag for flag in map(_stmt_frozen_flag, cls.body) if flag is not None) - return body_flags[-1] if body_flags else None + body_flags: Final = tuple(_stmt_frozen_flag(stmt, module_assignments) 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))) models: Final = _pydantic_classes(tree) bases_of: Final = {cls: _base_names(cls) for cls in models} - override_of: Final = {cls: _class_frozen_override(cls) for cls in models} + override_of: Final = {cls: _class_frozen_override(cls, module_assignments) for cls in models} def frozen(known: frozenset[str]) -> frozenset[str]: grown: Final = known | frozenset( diff --git a/scripts/type_discipline_gate.py b/scripts/type_discipline_gate.py index fda6bb5def0..3a0564c3928 100644 --- a/scripts/type_discipline_gate.py +++ b/scripts/type_discipline_gate.py @@ -203,7 +203,7 @@ def cmd_check(base: str) -> None: "`# pyright: ignore[rule] # `, `# mutable-ok: `, " "`# cast-ok: `, `# guard-ok: `, `# kwargs-ok: `, " "`# rebind-ok: `, `# writable-ok: `, " - "`# comprehension-ok: `), or remove an equal " + "`# comprehension-ok: `, `# frozen-ok: `), or remove an equal " "number elsewhere; the ceiling " "is the limit in type-discipline-budget.json." ) diff --git a/tests/unit/test_check_type_discipline.py b/tests/unit/test_check_type_discipline.py index 171e1460b0a..60852654776 100644 --- a/tests/unit/test_check_type_discipline.py +++ b/tests/unit/test_check_type_discipline.py @@ -599,6 +599,37 @@ def test_configdict_frozen_true_is_clean(tmp_path): assert "LIT015" not in _codes(tmp_path, src) +def test_shared_frozen_configdict_constant_is_clean(tmp_path): + src: Final = ( + "from typing import Final\n" + "from pydantic import BaseModel, ConfigDict\n" + "_RESPONSE_CONFIG: Final = ConfigDict(frozen=True)\n" + "class Foo(BaseModel):\n" + " model_config = _RESPONSE_CONFIG\n" + " x: int\n" + ) + assert "LIT015" not in _codes(tmp_path, src) + + +def test_last_shared_non_frozen_configdict_assignment_is_flagged(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" + ) + assert "LIT015" in _codes(tmp_path, src) + + +def test_unknown_model_config_name_is_flagged(tmp_path): + src: Final = ( + "from pydantic import BaseModel\nclass Foo(BaseModel):\n model_config = UNKNOWN_CONFIG\n x: int\n" + ) + assert "LIT015" in _codes(tmp_path, src) + + def test_dict_literal_model_config_frozen_true_is_clean(tmp_path): src = "from pydantic import BaseModel\nclass P(BaseModel):\n model_config = {'frozen': True, 'extra': 'allow'}\n" assert "LIT015" not in _codes(tmp_path, src) diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52b7682464b..aeb77212e8e 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -39,6 +39,6 @@ "limit": 369 }, "LIT015": { - "limit": 1450 + "limit": 1448 } }