From 5a4d8f9f04d890a4a7c8c92aa36b6f1b677e5336 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 8 Aug 2026 02:04:21 +0000 Subject: [PATCH] ci: resolve module-level constants when detecting extra="allow" Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../ban_pydantic_extra_allow.py | 90 +++++++++++-------- .../test_ban_pydantic_extra_allow.py | 32 +++++++ 2 files changed, 85 insertions(+), 37 deletions(-) diff --git a/tests/code_coverage_tests/ban_pydantic_extra_allow.py b/tests/code_coverage_tests/ban_pydantic_extra_allow.py index 8ab0d3d10c1..f9365ffea29 100644 --- a/tests/code_coverage_tests/ban_pydantic_extra_allow.py +++ b/tests/code_coverage_tests/ban_pydantic_extra_allow.py @@ -9,7 +9,8 @@ must declare its fields. import ast import os import sys -from typing import Final, Iterator, NamedTuple, Sequence +from types import MappingProxyType +from typing import Final, Iterator, Mapping, NamedTuple, Sequence SCAN_ROOT: Final = "litellm" @@ -78,67 +79,81 @@ class Violation(NamedTuple): return f"{self.file}::{self.model}" -def _is_extra_allow_keyword(keyword: ast.keyword) -> bool: - return keyword.arg == "extra" and isinstance(keyword.value, ast.Constant) and keyword.value.value == "allow" +def _assigned_names(statement: ast.stmt) -> tuple[tuple[str, ast.expr], ...]: + if isinstance(statement, ast.Assign): + targets, value = statement.targets, statement.value + elif isinstance(statement, ast.AnnAssign) and statement.value is not None: + targets, value = [statement.target], statement.value + else: + return () + return tuple((target.id, value) for target in targets if isinstance(target, ast.Name)) -def _mapping_sets_extra_allow(node: ast.Dict) -> bool: +def _module_constants(body: Sequence[ast.stmt]) -> Mapping[str, ast.expr]: + """Module-level ``NAME = value`` bindings, so a config held in a constant is still seen.""" + return MappingProxyType(dict(binding for statement in body for binding in _assigned_names(statement))) + + +def _resolve(node: ast.expr, constants: Mapping[str, ast.expr], seen: frozenset[str] = frozenset()) -> ast.expr: + if isinstance(node, ast.Name) and node.id in constants and node.id not in seen: + return _resolve(constants[node.id], constants, seen | {node.id}) + return node + + +def _is_allow_literal(node: ast.expr, constants: Mapping[str, ast.expr]) -> bool: + resolved = _resolve(node, constants) + if isinstance(resolved, ast.Constant): + return resolved.value == "allow" + return isinstance(resolved, ast.Attribute) and resolved.attr == "allow" + + +def _is_extra_allow_keyword(keyword: ast.keyword, constants: Mapping[str, ast.expr]) -> bool: + return keyword.arg == "extra" and _is_allow_literal(keyword.value, constants) + + +def _mapping_sets_extra_allow(node: ast.Dict, constants: Mapping[str, ast.expr]) -> bool: return any( - isinstance(key, ast.Constant) - and key.value == "extra" - and isinstance(value, ast.Constant) - and value.value == "allow" + isinstance(key, ast.Constant) and key.value == "extra" and _is_allow_literal(value, constants) for key, value in zip(node.keys, node.values) ) -def _is_extra_allow_value(node: ast.expr) -> bool: - if isinstance(node, ast.Call): - return any(_is_extra_allow_keyword(keyword) for keyword in node.keywords) - if isinstance(node, ast.Dict): - return _mapping_sets_extra_allow(node) +def _is_extra_allow_value(node: ast.expr, constants: Mapping[str, ast.expr]) -> bool: + resolved = _resolve(node, constants) + if isinstance(resolved, ast.Call): + return any(_is_extra_allow_keyword(keyword, constants) for keyword in resolved.keywords) + if isinstance(resolved, ast.Dict): + return _mapping_sets_extra_allow(resolved, constants) return False -def _assigns_extra_allow(statement: ast.stmt, target_names: Sequence[str]) -> bool: - if isinstance(statement, ast.Assign): - targets = statement.targets - value = statement.value - elif isinstance(statement, ast.AnnAssign) and statement.value is not None: - targets = [statement.target] - value = statement.value - else: - return False - names = {target.id for target in targets if isinstance(target, ast.Name)} - return bool(names & set(target_names)) and _is_extra_allow_value(value) +def _assigns_extra_allow(statement: ast.stmt, target_names: Sequence[str], constants: Mapping[str, ast.expr]) -> bool: + return any( + name in set(target_names) and _is_extra_allow_value(value, constants) + for name, value in _assigned_names(statement) + ) -def _is_allow_literal(node: ast.expr) -> bool: - if isinstance(node, ast.Constant): - return node.value == "allow" - return isinstance(node, ast.Attribute) and node.attr == "allow" - - -def _legacy_config_sets_extra_allow(class_def: ast.ClassDef) -> bool: +def _legacy_config_sets_extra_allow(class_def: ast.ClassDef, constants: Mapping[str, ast.expr]) -> bool: 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) + and _is_allow_literal(inner.value, constants) for inner in statement.body ) for statement in class_def.body ) -def _class_allows_extra(class_def: ast.ClassDef) -> bool: - if any(_is_extra_allow_keyword(keyword) for keyword in class_def.keywords): +def _class_allows_extra(class_def: ast.ClassDef, constants: Mapping[str, ast.expr]) -> bool: + if any(_is_extra_allow_keyword(keyword, constants) for keyword in class_def.keywords): return True - if any(_assigns_extra_allow(statement, ["model_config"]) for statement in class_def.body): + if any(_assigns_extra_allow(statement, ["model_config"], constants) for statement in class_def.body): return True - return _legacy_config_sets_extra_allow(class_def) + return _legacy_config_sets_extra_allow(class_def, constants) def _iter_classes(body: Sequence[ast.stmt], prefix: str = "") -> Iterator[tuple[str, ast.ClassDef]]: @@ -151,10 +166,11 @@ def _iter_classes(body: Sequence[ast.stmt], prefix: str = "") -> Iterator[tuple[ def find_violations_in_source(source: str, relative_path: str) -> tuple[Violation, ...]: tree = ast.parse(source, filename=relative_path) + constants = _module_constants(tree.body) return tuple( Violation(file=relative_path, line=class_def.lineno, model=qualified) for qualified, class_def in _iter_classes(tree.body) - if _class_allows_extra(class_def) + if _class_allows_extra(class_def, constants) ) diff --git a/tests/test_litellm/test_ban_pydantic_extra_allow.py b/tests/test_litellm/test_ban_pydantic_extra_allow.py index 4cab49f05bb..b6f050e1ce9 100644 --- a/tests/test_litellm/test_ban_pydantic_extra_allow.py +++ b/tests/test_litellm/test_ban_pydantic_extra_allow.py @@ -44,6 +44,26 @@ _REPO_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".." "class Foo(BaseModel):\n class Config:\n extra = Extra.allow\n", id="legacy_inner_config_enum", ), + pytest.param( + 'ALLOW = ConfigDict(extra="allow")\n\n\nclass Foo(BaseModel):\n model_config = ALLOW\n', + id="module_constant_config_dict", + ), + pytest.param( + 'ALLOW = {"extra": "allow"}\n\n\nclass Foo(BaseModel):\n model_config = ALLOW\n', + id="module_constant_dict", + ), + pytest.param( + 'ALLOW = {"extra": "allow"}\nALIAS = ALLOW\n\n\nclass Foo(BaseModel):\n model_config = ALIAS\n', + id="module_constant_chain", + ), + pytest.param( + 'EXTRA = "allow"\n\n\nclass Foo(BaseModel):\n model_config = ConfigDict(extra=EXTRA)\n', + id="module_constant_extra_value", + ), + pytest.param( + 'EXTRA = "allow"\n\n\nclass Foo(BaseModel, extra=EXTRA):\n pass\n', + id="module_constant_class_keyword", + ), ], ) def test_detects_extra_allow(source): @@ -74,6 +94,18 @@ def test_detects_extra_allow(source): 'def f():\n return ConfigDict(extra="allow")\n', id="outside_class", ), + pytest.param( + 'FORBID = ConfigDict(extra="forbid")\n\n\nclass Foo(BaseModel):\n model_config = FORBID\n', + id="module_constant_forbid", + ), + pytest.param( + "class Foo(BaseModel):\n model_config = IMPORTED_CONFIG\n", + id="unresolvable_constant", + ), + pytest.param( + "SELF = SELF\n\n\nclass Foo(BaseModel):\n model_config = SELF\n", + id="self_referencing_constant", + ), ], ) def test_ignores_non_violations(source):