diff --git a/tests/code_coverage_tests/ban_pydantic_extra_allow.py b/tests/code_coverage_tests/ban_pydantic_extra_allow.py index f9bd3a3cfcb..37c9d1d76b4 100644 --- a/tests/code_coverage_tests/ban_pydantic_extra_allow.py +++ b/tests/code_coverage_tests/ban_pydantic_extra_allow.py @@ -119,7 +119,7 @@ def _statements(node: ast.AST) -> tuple[ast.stmt, ...]: return tuple(statement for statement, _ in _scoped_statements(node)) -def _module_bindings(node: ast.AST) -> tuple[Binding, ...]: +def _bindings(node: ast.AST) -> tuple[Binding, ...]: return tuple( Binding(statement.lineno, name, value, branch) for statement, branch in _scoped_statements(node) @@ -166,26 +166,39 @@ def _assigns_extra_allow(statement: ast.stmt, target_names: Sequence[str], scope ) -def _legacy_config_sets_extra_allow(class_def: ast.ClassDef, scope: Scope) -> bool: +def _body_scope(node: ast.AST, enclosing: tuple[Binding, ...]) -> tuple[Binding, ...]: + """The bindings a statement in ``node``'s body reads, so a class-local constant is resolved + and shadows a module-level name of its own, the way the class body itself would read it.""" + local: Final = _bindings(node) + shadowed: Final = frozenset(binding.name for binding in local) + return tuple(binding for binding in enclosing if binding.name not in shadowed) + local + + +def _legacy_config_sets_extra_allow(config: ast.ClassDef, enclosing: tuple[Binding, ...]) -> bool: + bindings: Final = _body_scope(config, enclosing) return any( - isinstance(statement, ast.ClassDef) - and statement.name == "Config" - and any( - isinstance(inner, ast.Assign) - and any(isinstance(target, ast.Name) and target.id == "extra" for target in inner.targets) - and _is_allow_literal(inner.value, scope) - for inner in _statements(statement) - ) - for statement in _statements(class_def) + isinstance(inner, ast.Assign) + and any(isinstance(target, ast.Name) and target.id == "extra" for target in inner.targets) + and _is_allow_literal(inner.value, Scope(bindings, inner.lineno)) + for inner in _statements(config) ) -def _class_allows_extra(class_def: ast.ClassDef, scope: Scope) -> bool: - if any(_is_extra_allow_keyword(keyword, scope) for keyword in class_def.keywords): +def _class_allows_extra(class_def: ast.ClassDef, enclosing: tuple[Binding, ...]) -> bool: + if any(_is_extra_allow_keyword(keyword, Scope(enclosing, class_def.lineno)) for keyword in class_def.keywords): return True - if any(_assigns_extra_allow(statement, ["model_config"], scope) for statement in _statements(class_def)): + bindings: Final = _body_scope(class_def, enclosing) + statements: Final = _statements(class_def) + if any( + _assigns_extra_allow(statement, ["model_config"], Scope(bindings, statement.lineno)) for statement in statements + ): return True - return _legacy_config_sets_extra_allow(class_def, scope) + return any( + isinstance(statement, ast.ClassDef) + and statement.name == "Config" + and _legacy_config_sets_extra_allow(statement, bindings) + for statement in statements + ) def _iter_classes(node: ast.AST, prefix: str = "") -> Iterator[tuple[str, ast.ClassDef]]: @@ -203,11 +216,11 @@ def find_violations_in_source(source: str, relative_path: str) -> tuple[Violatio if "allow" not in source: return () tree: Final = ast.parse(source, filename=relative_path) - bindings: Final = _module_bindings(tree) + bindings: Final = _bindings(tree) return tuple( Violation(file=relative_path, line=class_def.lineno, model=qualified) for qualified, class_def in _iter_classes(tree) - if _class_allows_extra(class_def, Scope(bindings, class_def.lineno)) + if _class_allows_extra(class_def, bindings) ) diff --git a/tests/test_litellm/test_ban_pydantic_extra_allow.py b/tests/test_litellm/test_ban_pydantic_extra_allow.py index d302f558f1c..1bfaa3c07de 100644 --- a/tests/test_litellm/test_ban_pydantic_extra_allow.py +++ b/tests/test_litellm/test_ban_pydantic_extra_allow.py @@ -124,6 +124,19 @@ _ratchet_spec.loader.exec_module(ratchet) "class Foo(BaseModel):\n model_config = CONFIG\n", id="one_branch_of_a_conditional_binding_allows", ), + pytest.param( + 'class Foo(BaseModel):\n _CONFIG = ConfigDict(extra="allow")\n model_config = _CONFIG\n', + id="class_local_constant", + ), + pytest.param( + 'FORBID = ConfigDict(extra="forbid")\n\n\nclass Foo(BaseModel):\n' + ' FORBID = ConfigDict(extra="allow")\n model_config = FORBID\n', + id="class_local_constant_shadowing_a_harmless_module_one", + ), + pytest.param( + 'class Foo(BaseModel):\n class Config:\n _EXTRA = "allow"\n extra = _EXTRA\n', + id="legacy_inner_config_reading_its_own_constant", + ), ], ) def test_detects_extra_allow(source): @@ -181,6 +194,15 @@ def test_detects_extra_allow(source): "class Foo(BaseModel):\n model_config = CONFIG\n", id="conditional_allow_rebound_unconditionally_before_the_class", ), + pytest.param( + 'ALLOW = ConfigDict(extra="allow")\n\n\nclass Foo(BaseModel):\n' + ' ALLOW = ConfigDict(extra="forbid")\n model_config = ALLOW\n', + id="class_local_constant_shadowing_a_permissive_module_one", + ), + pytest.param( + 'class Foo(BaseModel):\n model_config = LATER\n LATER = ConfigDict(extra="allow")\n', + id="class_local_constant_bound_below_the_line_reading_it", + ), ], ) def test_ignores_non_violations(source):