Allow non-annotated assigns in template_fields (#20)

feluelle · Jul 4, 2022 · 44c1ac6 · 44c1ac6
1 parent e4549ac
commit 44c1ac6
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/airflint/rules/use_jinja_variable_get.py b/airflint/rules/use_jinja_variable_get.py
@@ -69,9 +69,15 @@ def _lookup_template_fields(self, keyword: ast.keyword) -> None:
         with open(file_path) as file:
             module = ast.parse(file.read())
         assert any(
-            isinstance(stmt, ast.AnnAssign)
-            and isinstance(stmt.target, ast.Name)
-            and stmt.target.id == "template_fields"
+            (
+                # E.g. <name>: Any = (..., )
+                isinstance(stmt, ast.AnnAssign)
+                and isinstance(target := stmt.target, ast.Name)
+                # E.g. <name> = (..., )
+                or isinstance(stmt, ast.Assign)
+                and isinstance(target := stmt.targets[0], ast.Name)
+            )
+            and target.id == "template_fields"
             and isinstance(stmt.value, ast.Tuple)
             and any(
                 isinstance(elt, ast.Constant) and elt.value == keyword.arg