Add support for annotated assignments to static attribute lookup.

karlotness · karlotness · commit 685c80c864f0 · 2022-06-18T22:24:36.000-04:00
When walking the ast of a module, look for AnnAssign nodes in addition
to Assign to support assignments with type annotations, for example.
Since we have to read different attributes, split the generators into
a for loop. Existing ast.Assign nodes follow the same processing as
before.
diff --git a/setuptools/config/expand.py b/setuptools/config/expand.py
@@ -69,24 +69,21 @@ def __init__(self, name: str, spec: ModuleSpec):
     def __getattr__(self, attr):
         """Attempt to load an attribute "statically", via :func:`ast.literal_eval`."""
         try:
-            assignment_expressions = (
-                statement
-                for statement in self.module.body
-                if isinstance(statement, ast.Assign)
-            )
-            expressions_with_target = (
-                (statement, target)
-                for statement in assignment_expressions
-                for target in statement.targets
-            )
-            matching_values = (
-                statement.value
-                for statement, target in expressions_with_target
-                if isinstance(target, ast.Name) and target.id == attr
-            )
-            return next(ast.literal_eval(value) for value in matching_values)
+            for statement in self.module.body:
+                if isinstance(statement, ast.Assign):
+                    targets = statement.targets
+                    value = statement.value
+                elif isinstance(statement, ast.AnnAssign):
+                    targets = [statement.target]
+                    value = statement.value
+                else:
+                    continue
+                for target in targets:
+                    if isinstance(target, ast.Name) and target.id == attr:
+                        return ast.literal_eval(value)
         except Exception as e:
             raise AttributeError(f"{self.name} has no attribute {attr}") from e
+        raise AttributeError(f"{self.name} has no attribute {attr}")
 
 
 def glob_relative(
diff --git a/setuptools/tests/config/test_expand.py b/setuptools/tests/config/test_expand.py
@@ -85,6 +85,18 @@ def test_read_attr(self, tmp_path, monkeypatch):
         values = expand.read_attr('lib.mod.VALUES', {'lib': 'pkg/sub'}, tmp_path)
         assert values['c'] == (0, 1, 1)
 
+    def test_read_annotated_attr(self, tmp_path):
+        files = {
+            "pkg/__init__.py": "",
+            "pkg/sub/__init__.py": (
+                "VERSION: str = '0.1.1'\n"
+                "raise SystemExit(1)\n"
+            ),
+        }
+        write_files(files, tmp_path)
+        # Make sure this attribute can be read statically
+        assert expand.read_attr('pkg.sub.VERSION', root_dir=tmp_path) == '0.1.1'
+
     def test_import_order(self, tmp_path):
         """
         Sometimes the import machinery will import the parent package of a nested
