Add support for parsing function type comments

This commit exposes two new attributes to the FunctionDef nodes, type_comment_args
respectively type_comment_annotations. These two attributes hold the type annotations
provided via type comments.

1 files changed, 32 insertions, 5 deletions

diff --git a/astroid/rebuilder.py b/astroid/rebuilder.py
index 223ddcdd..911d5601 100644
--- a/astroid/rebuilder.py
+++ b/astroid/rebuilder.py
@@ -13,7 +13,7 @@ order to get a single Astroid representation
 import sys
 
 import astroid
-from astroid._ast import _parse, _get_parser_module
+from astroid._ast import _parse, _get_parser_module, parse_function_type_comment
 from astroid import nodes
 
 
@@ -259,6 +259,24 @@ class TreeRebuilder(object):
 
         return type_object.value
 
+    def check_function_type_comment(self, node):
+        type_comment = getattr(node, 'type_comment', None)
+        if not type_comment:
+            return None
+
+        try:
+            type_comment_ast = parse_function_type_comment(type_comment)
+        except SyntaxError:
+            # Invalid type comment, just skip it.
+            return None
+
+        returns = None
+        argtypes = [self.visit(elem, node) for elem in (type_comment_ast.argtypes or [])]
+        if type_comment_ast.returns:
+            returns = self.visit(type_comment_ast.returns, node)
+
+        return returns, argtypes
+
     def visit_assign(self, node, parent):
         """visit a Assign node by returning a fresh instance of it"""
         type_annotation = self.check_type_comment(node)
@@ -528,10 +546,19 @@ class TreeRebuilder(object):
             returns = self.visit(node.returns, newnode)
         else:
             returns = None
-        newnode.postinit(self.visit(node.args, newnode),
-                         [self.visit(child, newnode)
-                          for child in node.body],
-                         decorators, returns)
+
+        type_comment_args = type_comment_returns = None
+        type_comment_annotation = self.check_function_type_comment(node)
+        if type_comment_annotation:
+            type_comment_returns, type_comment_args = type_comment_annotation
+        newnode.postinit(
+            args=self.visit(node.args, newnode),
+            body=[self.visit(child, newnode) for child in node.body],
+            decorators=decorators,
+            returns=returns,
+            type_comment_returns=type_comment_returns,
+            type_comment_args=type_comment_args,
+        )
         self._global_names.pop()
         return newnode