diff --git a/scripts/check_type_discipline.py b/scripts/check_type_discipline.py index d210c7849d8..a9f417f807c 100644 --- a/scripts/check_type_discipline.py +++ b/scripts/check_type_discipline.py @@ -865,14 +865,14 @@ def iter_typeddict_violations(path: Path, tree: ast.AST) -> Iterator[Violation]: def _pydantic_classes(tree: ast.AST) -> tuple[ast.ClassDef, ...]: classes: Final = tuple(node for node in ast.walk(tree) if isinstance(node, ast.ClassDef)) - bases_of: Final = {cls.name: _base_names(cls) for cls in classes} def expand(known: frozenset[str]) -> frozenset[str]: - grown: Final = known | frozenset(name for name, bases in bases_of.items() if bases & known) + grown: Final = known | frozenset(cls.name for cls in classes if _base_names(cls) & known) return grown if grown == known else expand(grown) - names: Final = expand(PYDANTIC_BASES) - expand(frozenset((TYPEDDICT_BASE,))) - return tuple(cls for cls in classes if cls.name in names) + model_names: Final = expand(PYDANTIC_BASES) + typeddict_names: Final = expand(frozenset((TYPEDDICT_BASE,))) + return tuple(cls for cls in classes if _base_names(cls) & model_names and not _base_names(cls) & typeddict_names) def _bool_constant(value: ast.expr) -> bool | None: diff --git a/tests/unit/test_check_type_discipline.py b/tests/unit/test_check_type_discipline.py index 9ea4508ebc3..bdfb909833f 100644 --- a/tests/unit/test_check_type_discipline.py +++ b/tests/unit/test_check_type_discipline.py @@ -13,6 +13,7 @@ import re import subprocess import sys from pathlib import Path +from typing import Final import pytest @@ -770,6 +771,21 @@ def test_same_named_models_use_each_classes_own_override(tmp_path): assert _codes(tmp_path, src) == ["LIT015"] +def test_pydantic_model_discovery_uses_each_class_nodes_bases(tmp_path): + src: Final = ( + "from pydantic import BaseModel\n" + "class Config(BaseModel):\n" + " x: int\n" + "class Consumer(BaseModel, frozen=True):\n" + " class Config:\n" + " arbitrary_types_allowed = True\n" + ) + path: Final = tmp_path / "snippet.py" + path.write_text(src, encoding="utf-8") + violations: Final = checker.check_file(path) + assert [(violation.line, violation.code) for violation in violations] == [(2, "LIT015")] + + # --------------------------------------------------------------------------- # # Stacked comprehension clauses (LIT014) # --------------------------------------------------------------------------- #