Add support for type comments (#548)

ChangeLog

@@ -2,6 +2,12 @@ Change log for the astroid package (used to be astng)
 =====================================================
 
 -- 2.0
+--
+   * Switched to using typed_ast for getting access to type comments
+
+     As a side effect of this change, some nodes gained a new `type_annotation` attribute,
+     which, if the type comments were correctly parsed, should contain a node object
+     with the corresponding objects from the type comment.
 
    * typing.X[...] and typing.NewType are inferred as classes instead of instances.
 

astroid/__pkginfo__.py

@@ -6,6 +6,7 @@
 # For details: https://github.com/PyCQA/astroid/blob/master/COPYING.LESSER
 
 """astroid packaging information"""
+import platform
 
 distname = 'astroid'
 
@@ -19,9 +20,12 @@
     'lazy_object_proxy',
     'six',
     'wrapt',
-    'typing;python_version<"3.5"'
+    'typing;python_version<"3.5"',
 ]
 
+if platform.python_implementation() == 'CPython':
+    install_requires.append('typed_ast;python_version<"3.7"')
+
 # pylint: disable=redefined-builtin; why license is a builtin anyway?
 license = 'LGPL'
 

astroid/_ast.py

@@ -0,0 +1,21 @@
+import ast
+
+_ast_py2 = _ast_py3 = None
+try:
+    import typed_ast.ast3 as _ast_py3
+    import typed_ast.ast27 as _ast_py2
+except ImportError:
+    pass
+
+
+def _get_parser_module(parse_python_two: bool = False):
+    if parse_python_two:
+        parser_module = _ast_py2
+    else:
+        parser_module = _ast_py3
+    return parser_module or ast
+
+
+def _parse(string: str,
+           parse_python_two: bool = False):
+    return _get_parser_module(parse_python_two=parse_python_two).parse(string)

astroid/builder.py

@@ -16,8 +16,9 @@
 import os
 import sys
 import textwrap
-import _ast
 
+
+from astroid._ast import _parse
 from astroid import bases
 from astroid import exceptions
 from astroid import manager
@@ -37,10 +38,6 @@
 _STATEMENT_SELECTOR = '#@'
 
 
-def _parse(string):
-    return compile(string, "<string>", 'exec', _ast.PyCF_ONLY_AST)
-
-
 if sys.version_info >= (3, 0):
     from tokenize import detect_encoding
 

astroid/node_classes.py

@@ -1699,6 +1699,7 @@ class Assign(mixins.AssignTypeMixin, Statement):
     <Assign l.1 at 0x7effe1db8550>
     """
     _astroid_fields = ('targets', 'value',)
+    _other_other_fields = ('type_annotation',)
     targets = None
     """What is being assigned to.
 
@@ -1709,8 +1710,13 @@ class Assign(mixins.AssignTypeMixin, Statement):
 
     :type: NodeNG or None
     """
+    type_annotation = None
+    """If present, this will contain the type annotation passed by a type comment
 
-    def postinit(self, targets=None, value=None):
+    :type: NodeNG or None
+    """
+
+    def postinit(self, targets=None, value=None, type_annotation=None):
         """Do some setup after initialisation.
 
         :param targets: What is being assigned to.
@@ -1721,6 +1727,7 @@ def postinit(self, targets=None, value=None):
         """
         self.targets = targets
         self.value = value
+        self.type_annotation = type_annotation
 
     def get_children(self):
         yield from self.targets
@@ -2912,6 +2919,7 @@ class For(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn,
     <For l.1 at 0x7f23b2e8cf28>
     """
     _astroid_fields = ('target', 'iter', 'body', 'orelse',)
+    _other_other_fields = ('type_annotation',)
     _multi_line_block_fields = ('body', 'orelse')
     target = None
     """What the loop assigns to.
@@ -2933,9 +2941,14 @@ class For(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn,
 
     :type: list(NodeNG) or None
     """
+    type_annotation = None
+    """If present, this will contain the type annotation passed by a type comment
+
+    :type: NodeNG or None
+    """
 
     # pylint: disable=redefined-builtin; had to use the same name as builtin ast module.
-    def postinit(self, target=None, iter=None, body=None, orelse=None):
+    def postinit(self, target=None, iter=None, body=None, orelse=None, type_annotation=None):
         """Do some setup after initialisation.
 
         :param target: What the loop assigns to.
@@ -2954,6 +2967,7 @@ def postinit(self, target=None, iter=None, body=None, orelse=None):
         self.iter = iter
         self.body = body
         self.orelse = orelse
+        self.type_annotation = type_annotation
 
     optional_assign = True
     """Whether this node optionally assigns a variable.
@@ -4231,7 +4245,8 @@ class With(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn,
     >>> node
     <With l.2 at 0x7f23b2e4e710>
     """
-    _astroid_fields = ('items', 'body')
+    _astroid_fields = ('items', 'body',)
+    _other_other_fields = ('type_annotation',)
     _multi_line_block_fields = ('body',)
     items = None
     """The pairs of context managers and the names they are assigned to.
@@ -4243,8 +4258,13 @@ class With(mixins.MultiLineBlockMixin, mixins.BlockRangeMixIn,
 
     :type: list(NodeNG) or None
     """
+    type_annotation = None
+    """If present, this will contain the type annotation passed by a type comment
+
+    :type: NodeNG or None
+    """
 
-    def postinit(self, items=None, body=None):
+    def postinit(self, items=None, body=None, type_annotation=None):
         """Do some setup after initialisation.
 
         :param items: The pairs of context managers and the names
@@ -4256,6 +4276,7 @@ def postinit(self, items=None, body=None):
         """
         self.items = items
         self.body = body
+        self.type_annotation = type_annotation
 
     @decorators.cachedproperty
     def blockstart_tolineno(self):