Expose function annotation to astroid. `Arguments` node exposes 'varargannotation', 'kwargannotation' and 'annotations' attributes, while `Function` node has the 'returns' attribute.

5 files changed, 99 insertions, 12 deletions

diff --git a/ChangeLog b/ChangeLog
index 08fd892..ab7aec0 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -13,6 +13,10 @@ Change log for the astroid package (used to be astng)
     * Add `slots` method to `Class` nodes, for retrieving
       the list of valid slots it defines.
 
+    * Expose function annotation to astroid: `Arguments` node
+      exposes 'varargannotation', 'kwargannotation' and 'annotations'
+      attributes, while `Function` node has the 'returns' attribute.
+
 2014-04-30  --  1.1.1
     * `Class.metaclass()` looks in ancestors when the current class
       does not define explicitly a metaclass.
diff --git a/node_classes.py b/node_classes.py
index 01dc8d9..5a70e7e 100644
--- a/node_classes.py
+++ b/node_classes.py
@@ -26,6 +26,8 @@ from astroid.bases import (NodeNG, Statement, Instance, InferenceContext,
 from astroid.mixins import BlockRangeMixIn, AssignTypeMixin, \
                                  ParentAssignTypeMixin, FromImportMixIn
 
+PY3K = sys.version_info >= (3, 0)
+
 
 def unpack_infer(stmt, context=None):
     """recursively generate nodes inferred by the given statement.
@@ -253,7 +255,26 @@ class Name(LookupMixIn, NodeNG):
 
 class Arguments(NodeNG, AssignTypeMixin):
     """class representing an Arguments node"""
-    _astroid_fields = ('args', 'defaults', 'kwonlyargs', 'kw_defaults')
+    if PY3K:
+        # Python 3.4+ uses a different approach regarding annotations,
+        # each argument is a new class, _ast.arg, which exposes an
+        # 'annotation' attribute. In astroid though, arguments are exposed
+        # as is in the Arguments node and the only way to expose annotations
+        # is by using something similar with Python 3.3:
+        #  - we expose 'varargannotation' and 'kwargannotation' of annotations
+        #    of varargs and kwargs.
+        #  - we expose 'annotation', a list with annotations for
+        #    for each normal argument. If an argument doesn't have an
+        #    annotation, its value will be None.
+
+        _astroid_fields = ('args', 'defaults', 'kwonlyargs',
+                           'kw_defaults', 'annotations',
+                           'varargannotation', 'kwargannotation')
+        annotations = None
+        varargannotation = None
+        kwargannotation = None
+    else:
+        _astroid_fields = ('args', 'defaults', 'kwonlyargs', 'kw_defaults')
     args = None
     defaults = None
     kwonlyargs = None
diff --git a/rebuilder.py b/rebuilder.py
index 47eff50..34b7b11 100644
--- a/rebuilder.py
+++ b/rebuilder.py
@@ -179,10 +179,26 @@ class TreeRebuilder(object):
         vararg, kwarg = node.vararg, node.kwarg
         # change added in 82732 (7c5c678e4164), vararg and kwarg
         # are instances of `_ast.arg`, not strings
-        if vararg and PY34:
-            vararg = vararg.arg
-        if kwarg and PY34:
-            kwarg = kwarg.arg
+        if vararg:
+            annotation = None
+            if PY34:
+                if vararg.annotation:
+                    annotation = self.visit(vararg.annotation, newnode)
+                vararg = vararg.arg
+            elif PY3K:
+                if node.varargannotation:
+                    annotation = self.visit(node.varargannotation, newnode)
+            newnode.varargannotation = annotation
+        if kwarg:
+            annotation = None
+            if PY34:
+                if kwarg.annotation:
+                    annotation = self.visit(kwarg.annotation, newnode)
+                kwarg = kwarg.arg
+            elif PY3K:
+                if node.kwargannotation:
+                    annotation = self.visit(node.kwargannotation, newnode)
+            newnode.kwargannotation = annotation
         newnode.vararg = vararg
         newnode.kwarg = kwarg
         # save argument names in locals:
@@ -492,6 +508,8 @@ class TreeRebuilder(object):
         decorators = getattr(node, attr)
         if decorators:
             newnode.decorators = self.visit_decorators(node, newnode)
+        if PY3K and node.returns:
+            newnode.returns = self.visit(node.returns, newnode)
         newnode.set_line_info(newnode.last_child())
         self._global_names.pop()
         frame = newnode.parent.frame()
@@ -831,6 +849,9 @@ class TreeRebuilder3k(TreeRebuilder):
         newnode.kwonlyargs = [self.visit(child, newnode) for child in node.kwonlyargs]
         self.asscontext = None
         newnode.kw_defaults = [self.visit(child, newnode) if child else None for child in node.kw_defaults]
+        newnode.annotations = [
+            self.visit(arg.annotation, newnode) if arg.annotation else None
+            for arg in node.args]
         return newnode
 
     def visit_excepthandler(self, node, parent):
diff --git a/scoped_nodes.py b/scoped_nodes.py
index 95d45a5..f070837 100644
--- a/scoped_nodes.py
+++ b/scoped_nodes.py
@@ -46,6 +46,7 @@ from astroid.bases import Statement
 from astroid.manager import AstroidManager
 
 ITER_METHODS = ('__iter__', '__getitem__')
+PY3K = sys.version_info >= (3, 0)
 
 
 def remove_nodes(func, cls):
@@ -562,7 +563,11 @@ class Lambda(LocalsDictNodeNG, FilterStmtsMixin):
 
 
 class Function(Statement, Lambda):
-    _astroid_fields = ('decorators', 'args', 'body')
+    if PY3K:
+        _astroid_fields = ('decorators', 'args', 'body', 'returns')
+        returns = None
+    else:
+        _astroid_fields = ('decorators', 'args', 'body')
 
     special_attributes = set(('__name__', '__doc__', '__dict__'))
     is_function = True
diff --git a/test/unittest_python3.py b/test/unittest_python3.py
index 295984a..f6d1453 100644
--- a/test/unittest_python3.py
+++ b/test/unittest_python3.py
@@ -15,12 +15,12 @@
 #
 # You should have received a copy of the GNU Lesser General Public License along
 # with astroid. If not, see <http://www.gnu.org/licenses/>.
-import sys
+
 from textwrap import dedent
 
 from logilab.common.testlib import TestCase, unittest_main, require_version
 
-from astroid.node_classes import Assign, Discard, YieldFrom
+from astroid.node_classes import Assign, Discard, YieldFrom, Name, Const
 from astroid.manager import AstroidManager
 from astroid.builder import AstroidBuilder
 from astroid.scoped_nodes import Class, Function
@@ -97,7 +97,7 @@ class Python3TC(TestCase):
         self.assertFalse(klass.metaclass())
 
     @require_version('3.0')
-    def test_metaclass_imported(self): 
+    def test_metaclass_imported(self):
         astroid = self.builder.string_build(dedent("""
         from abc import ABCMeta 
         class Test(metaclass=ABCMeta): pass"""))
@@ -105,7 +105,7 @@ class Python3TC(TestCase):
 
         metaclass = klass.metaclass()
         self.assertIsInstance(metaclass, Class)
-        self.assertEqual(metaclass.name, 'ABCMeta')       
+        self.assertEqual(metaclass.name, 'ABCMeta')
 
     @require_version('3.0')
     def test_as_string(self):
@@ -115,7 +115,7 @@ class Python3TC(TestCase):
         astroid = self.builder.string_build(body)
         klass = astroid.body[1]
 
-        self.assertEqual(klass.as_string(), 
+        self.assertEqual(klass.as_string(),
                          '\n\nclass Test(metaclass=ABCMeta):\n    pass\n')
 
     @require_version('3.0')
@@ -180,6 +180,42 @@ class Python3TC(TestCase):
                 meta = impl.metaclass()
                 self.assertIsInstance(meta, Class)
                 self.assertEqual(meta.name, metaclass)
-                         
+
+    @require_version('3.0')
+    def test_annotation_support(self):
+        astroid = self.builder.string_build(dedent("""
+        def test(a: int, b: str, c: None, d, e,
+                 *args: float, **kwargs: int)->int:
+            pass
+        """))
+        func = astroid['test']
+        self.assertIsInstance(func.args.varargannotation, Name)
+        self.assertEqual(func.args.varargannotation.name, 'float')
+        self.assertIsInstance(func.args.kwargannotation, Name)
+        self.assertEqual(func.args.kwargannotation.name, 'int')
+        self.assertIsInstance(func.returns, Name)
+        self.assertEqual(func.returns.name, 'int')
+        arguments = func.args
+        self.assertIsInstance(arguments.annotations[0], Name)
+        self.assertEqual(arguments.annotations[0].name, 'int')
+        self.assertIsInstance(arguments.annotations[1], Name)
+        self.assertEqual(arguments.annotations[1].name, 'str')
+        self.assertIsInstance(arguments.annotations[2], Const)
+        self.assertIsNone(arguments.annotations[2].value)
+        self.assertIsNone(arguments.annotations[3])
+        self.assertIsNone(arguments.annotations[4])
+
+        astroid = self.builder.string_build(dedent("""
+        def test(a: int=1, b: str=2):
+            pass
+        """))
+        func = astroid['test']
+        self.assertIsInstance(func.args.annotations[0], Name)
+        self.assertEqual(func.args.annotations[0].name, 'int')
+        self.assertIsInstance(func.args.annotations[1], Name)
+        self.assertEqual(func.args.annotations[1].name, 'str')
+        self.assertIsNone(func.returns)
+
+
 if __name__ == '__main__':
     unittest_main()