fixup! fixup! fixup! style: fix provider mypy error

Author: Lee-W

scripts/ci/prek/check_template_context_variable_in_sync.py

@@ -75,26 +75,76 @@ def extract_keys_from_dict(node: ast.Dict) -> typing.Iterator[str]:
                 raise ValueError("Key in dictionary is not a string literal")
             yield key.value
 
-    # Extract keys from the main `context` dictionary assignment
-    context_assignment: ast.AnnAssign = next(
-        stmt
-        for stmt in fn_get_template_context.body
-        if isinstance(stmt, ast.AnnAssign)
-        and isinstance(stmt.target, ast.Attribute)
-        and isinstance(stmt.target.value, ast.Name)
-        and stmt.target.value.id == "self"
-        and stmt.target.attr == "_cached_template_context"
-    )
-
-    if not isinstance(context_assignment.value, ast.BoolOp):
-        raise TypeError("Expected a BoolOp like 'self._cached_template_context or {...}'.")
+    # Extract keys from the main `context` dictionary assignment.
+    # Supports two patterns:
+    #   1. Annotated assignment: self._cached_template_context: T = self._cached_template_context or {...}
+    #   2. Guard block: if self._cached_template_context is None: self._cached_template_context = {...}
+    context_dict: ast.Dict | None = None
 
-    context_assignment_op = context_assignment.value
-    _, context_assignment_value = context_assignment_op.values
-
-    if not isinstance(context_assignment_value, ast.Dict):
-        raise ValueError("'context' is not assigned a dictionary literal")
-    yield from extract_keys_from_dict(context_assignment_value)
+    for stmt in fn_get_template_context.body:
+        # Pattern 1: AnnAssign with BoolOp
+        if (
+            isinstance(stmt, ast.AnnAssign)
+            and isinstance(stmt.target, ast.Attribute)
+            and isinstance(stmt.target.value, ast.Name)
+            and stmt.target.value.id == "self"
+            and stmt.target.attr == "_cached_template_context"
+            and isinstance(stmt.value, ast.BoolOp)
+        ):
+            _, context_dict = stmt.value.values  # type: ignore[assignment]
+            break
+        # Pattern 2: if self._cached_template_context is None: self._cached_template_context = {...}
+        if (
+            isinstance(stmt, ast.If)
+            and isinstance(stmt.test, ast.Compare)
+            and isinstance(stmt.test.left, ast.Attribute)
+            and isinstance(stmt.test.left.value, ast.Name)
+            and stmt.test.left.value.id == "self"
+            and stmt.test.left.attr == "_cached_template_context"
+            and len(stmt.test.ops) == 1
+            and isinstance(stmt.test.ops[0], ast.Is)
+            and len(stmt.test.comparators) == 1
+            and isinstance(stmt.test.comparators[0], ast.Constant)
+            and stmt.test.comparators[0].value is None
+        ):
+            inner = next(
+                (
+                    s
+                    for s in stmt.body
+                    if isinstance(s, (ast.Assign, ast.AnnAssign))
+                    and (
+                        (
+                            isinstance(s, ast.Assign)
+                            and any(
+                                isinstance(t, ast.Attribute)
+                                and isinstance(t.value, ast.Name)
+                                and t.value.id == "self"
+                                and t.attr == "_cached_template_context"
+                                for t in s.targets
+                            )
+                        )
+                        or (
+                            isinstance(s, ast.AnnAssign)
+                            and isinstance(s.target, ast.Attribute)
+                            and isinstance(s.target.value, ast.Name)
+                            and s.target.value.id == "self"
+                            and s.target.attr == "_cached_template_context"
+                        )
+                    )
+                ),
+                None,
+            )
+            if inner is not None:
+                val = inner.value if isinstance(inner, ast.AnnAssign) else inner.value  # type: ignore[union-attr]
+                if isinstance(val, ast.Dict):
+                    context_dict = val
+            break
+
+    if context_dict is None:
+        raise ValueError("Could not find '_cached_template_context' dictionary assignment in get_template_context()")
+    if not isinstance(context_dict, ast.Dict):
+        raise ValueError("'_cached_template_context' is not assigned a dictionary literal")
+    yield from extract_keys_from_dict(context_dict)
 
     # Handle keys added conditionally in `if from_server`
     for stmt in fn_get_template_context.body: