Support AnnAssign in StmtVisitor

Assignments with type annotations were added in python 3.6.

Author: bcaller

examples/example_inputs/assignment_with_annotation.py

@@ -0,0 +1,2 @@
+x: int
+y: int=5

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

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')