Merge pull request #143 from bcaller/bc-annassign

KevinHock · web-flow · commit a2c46158ffbe · 2018-07-12T18:50:16.000-07:00
Support AnnAssign in StmtVisitor
diff --git a/.codeclimate.yml b/.codeclimate.yml
@@ -23,7 +23,7 @@ exclude_paths:
 - "pyt/intraprocedural_cfg.py"
 - "pyt/repo_runner.py"
 - "pyt/save.py"
-- "example/**"
+- "examples/**"
 - "profiling/**"
 - "tests/**"
 - "LICENSE"
diff --git a/examples/example_inputs/assignment_with_annotation.py b/examples/example_inputs/assignment_with_annotation.py
@@ -0,0 +1,2 @@
+x: int
+y: int=5
diff --git a/pyt/cfg/stmt_visitor.py b/pyt/cfg/stmt_visitor.py
@@ -422,6 +422,14 @@ def visit_Assign(self, node):
                     path=self.filenames[-1]
                 ))
 
+    def visit_AnnAssign(self, node):
+        if node.value is None:
+            return IgnoredNode()
+        else:
+            assign = ast.Assign(targets=[node.target], value=node.value)
+            ast.copy_location(assign, node)
+            return self.visit(assign)
+
     def assignment_call_node(self, left_hand_label, ast_node):
         """Handle assignments that contain a function call on its right side."""
         self.undecided = True  # Used for handling functions in assignments
diff --git a/tests/cfg/cfg_test.py b/tests/cfg/cfg_test.py
@@ -720,6 +720,25 @@ def test_assignment_and_builtin_line_numbers(self):
         self.assertLineNumber(assign, 1)
         self.assertLineNumber(builtin, 2)
 
+    def test_assignment_with_annotation(self):
+        self.cfg_create_from_file('examples/example_inputs/assignment_with_annotation.py')
+
+        self.assert_length(self.cfg.nodes, expected_length=3)
+
+        entry = 0
+        assign = 1
+        exit_node = 2
+
+        self.assertInCfg([(assign, entry), (exit_node, assign)])
+        self.assertEqual(self.cfg.nodes[assign].label, 'y = 5')
+
+    def test_assignment_with_annotation_line_numbers(self):
+        self.cfg_create_from_file('examples/example_inputs/assignment_with_annotation.py')
+
+        assign = self.cfg.nodes[1]
+
+        self.assertLineNumber(assign, 2)
+
     def test_multiple_assignment(self):
         self.cfg_create_from_file('examples/example_inputs/assignment_multiple_assign.py')
 
