Can access per argument type comments (#667)

Close #665

4 files changed, 61 insertions, 0 deletions

diff --git a/ChangeLog b/ChangeLog
index 1a4e4eff..5b21c00b 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -51,6 +51,10 @@ Release Date: TBA
   Close PyCQA/pylint#2734
   Close PyCQA/pylint#2740
 
+* Can access per argument type comments through new ``Arguments.type_comment_args`` attribute.
+
+  Close #665
+
 
 What's New in astroid 2.2.0?
 ============================
diff --git a/astroid/node_classes.py b/astroid/node_classes.py
index 449470d0..2e39b060 100644
--- a/astroid/node_classes.py
+++ b/astroid/node_classes.py
@@ -1425,6 +1425,7 @@ class Arguments(mixins.AssignTypeMixin, NodeNG):
         "varargannotation",
         "kwargannotation",
         "kwonlyargs_annotations",
+        "type_comment_args",
     )
     varargannotation = None
     """The type annotation for the variable length arguments.
@@ -1499,6 +1500,15 @@ class Arguments(mixins.AssignTypeMixin, NodeNG):
         :type: list(NodeNG)
         """
 
+        self.type_comment_args = []
+        """The type annotation, passed by a type comment, of each argument.
+
+        If an argument does not have a type comment,
+        the value for that argument will be None.
+
+        :type: list(NodeNG or None)
+        """
+
     def postinit(
         self,
         args,
@@ -1509,6 +1519,7 @@ class Arguments(mixins.AssignTypeMixin, NodeNG):
         kwonlyargs_annotations=None,
         varargannotation=None,
         kwargannotation=None,
+        type_comment_args=None,
     ):
         """Do some setup after initialisation.
 
@@ -1543,6 +1554,10 @@ class Arguments(mixins.AssignTypeMixin, NodeNG):
         :param kwargannotation: The type annotation for the variable length
             keyword arguments.
         :type kwargannotation: NodeNG
+
+        :param type_comment_args: The type annotation,
+            passed by a type comment, of each argument.
+        :type type_comment_args: list(NodeNG or None)
         """
         self.args = args
         self.defaults = defaults
@@ -1552,6 +1567,7 @@ class Arguments(mixins.AssignTypeMixin, NodeNG):
         self.kwonlyargs_annotations = kwonlyargs_annotations
         self.varargannotation = varargannotation
         self.kwargannotation = kwargannotation
+        self.type_comment_args = type_comment_args
 
     def _infer_name(self, frame, name):
         if self.parent is frame:
diff --git a/astroid/rebuilder.py b/astroid/rebuilder.py
index 1b313ee2..c18893d8 100644
--- a/astroid/rebuilder.py
+++ b/astroid/rebuilder.py
@@ -220,6 +220,7 @@ class TreeRebuilder:
             kw_defaults = []
             annotations = []
             kwonlyargs_annotations = []
+        type_comment_args = [self.check_type_comment(child) for child in node.args]
 
         newnode.postinit(
             args=args,
@@ -230,6 +231,7 @@ class TreeRebuilder:
             kwonlyargs_annotations=kwonlyargs_annotations,
             varargannotation=varargannotation,
             kwargannotation=kwargannotation,
+            type_comment_args=type_comment_args,
         )
         # save argument names in locals:
         if vararg:
diff --git a/astroid/tests/unittest_nodes.py b/astroid/tests/unittest_nodes.py
index 7d747aff..800eef3a 100644
--- a/astroid/tests/unittest_nodes.py
+++ b/astroid/tests/unittest_nodes.py
@@ -1023,6 +1023,45 @@ def test_type_comments_function():
         assert node.type_comment_returns.as_string() == expected_returns_string
 
 
+@pytest.mark.skipif(not HAS_TYPED_AST, reason="requires typed_ast")
+def test_type_comments_arguments():
+    module = builder.parse(
+        """
+    def func(
+        a,  # type: int
+    ):
+        # type: (...) -> str
+        pass
+    def func1(
+        a,  # type: int
+        b,  # type: int
+        c,  # type: int
+    ):
+        # type: (...) -> (str, str)
+        pass
+    def func2(
+        a,  # type: int
+        b,  # type: int
+        c,  # type: str
+        d,  # type: List[int]
+    ):
+        # type: (...) -> List[int]
+        pass
+    """
+    )
+    expected_annotations = [
+        ["int"],
+        ["int", "int", "int"],
+        ["int", "int", "str", "List[int]"],
+    ]
+    for node, expected_args in zip(module.body, expected_annotations):
+        assert len(node.type_comment_args) == 1
+        assert isinstance(node.type_comment_args[0], astroid.Ellipsis)
+        assert len(node.args.type_comment_args) == len(expected_args)
+        for expected_arg, actual_arg in zip(expected_args, node.args.type_comment_args):
+            assert actual_arg.as_string() == expected_arg
+
+
 def test_is_generator_for_yield_assignments():
     node = astroid.extract_node(
         """