ci: resolve module-level constants when detecting extra="allow"

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

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)
     )
 
 

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):