Fix annotated assign statement

CQCL · Nov 27, 2023 · a9238ec · a9238ec
1 parent a3ba236
commit a9238ec
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 0 deletions.
diff --git a/guppy/cfg/bb.py b/guppy/cfg/bb.py
@@ -103,6 +103,11 @@ def visit_AugAssign(self, node: ast.AugAssign) -> None:
         for name in name_nodes_in_ast(node.target):
             self.stats.assigned[name.id] = node
 
+    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
+        self.stats.update_used(node.value)
+        for name in name_nodes_in_ast(node.target):
+            self.stats.assigned[name.id] = node
+
     def visit_NestedFunctionDef(self, node: NestedFunctionDef) -> None:
         # In order to compute the used external variables in a nested function
         # definition, we have to run live variable analysis first

diff --git a/guppy/cfg/builder.py b/guppy/cfg/builder.py
@@ -96,6 +96,11 @@ def visit_AugAssign(
     ) -> Optional[BB]:
         return self._build_node_value(node, bb)
 
+    def visit_AnnAssign(
+        self, node: ast.AnnAssign, bb: BB, jumps: Jumps
+    ) -> Optional[BB]:
+        return self._build_node_value(node, bb)
+
     def visit_Expr(self, node: ast.Expr, bb: BB, jumps: Jumps) -> Optional[BB]:
         # This is an expression statement where the value is discarded
         node.value, bb = ExprBuilder.build(node.value, self.cfg, bb)