Skip to content

Commit

Permalink
Merge d5248a7 into e0a298d
Browse files Browse the repository at this point in the history
  • Loading branch information
AWhetter committed May 1, 2019
2 parents e0a298d + d5248a7 commit ff55cf4
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 0 deletions.
4 changes: 4 additions & 0 deletions ChangeLog
Expand Up @@ -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?
============================
Expand Down
16 changes: 16 additions & 0 deletions astroid/node_classes.py
Expand Up @@ -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.
Expand Down Expand Up @@ -1499,6 +1500,15 @@ def __init__(self, vararg=None, kwarg=None, parent=None):
: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,
Expand All @@ -1509,6 +1519,7 @@ def postinit(
kwonlyargs_annotations=None,
varargannotation=None,
kwargannotation=None,
type_comment_args=None,
):
"""Do some setup after initialisation.
Expand Down Expand Up @@ -1543,6 +1554,10 @@ def postinit(
: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
Expand All @@ -1552,6 +1567,7 @@ def postinit(
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:
Expand Down
2 changes: 2 additions & 0 deletions astroid/rebuilder.py
Expand Up @@ -220,6 +220,7 @@ def visit_arguments(self, node, parent):
kw_defaults = []
annotations = []
kwonlyargs_annotations = []
type_comment_args = [self.check_type_comment(child) for child in node.args]

newnode.postinit(
args=args,
Expand All @@ -230,6 +231,7 @@ def visit_arguments(self, node, parent):
kwonlyargs_annotations=kwonlyargs_annotations,
varargannotation=varargannotation,
kwargannotation=kwargannotation,
type_comment_args=type_comment_args,
)
# save argument names in locals:
if vararg:
Expand Down
39 changes: 39 additions & 0 deletions astroid/tests/unittest_nodes.py
Expand Up @@ -1023,6 +1023,45 @@ def func2():
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(
"""
Expand Down

0 comments on commit ff55cf4

Please sign in to comment.