diff options
| -rw-r--r-- | ChangeLog | 4 | ||||
| -rw-r--r-- | astroid/node_classes.py | 16 | ||||
| -rw-r--r-- | astroid/rebuilder.py | 2 | ||||
| -rw-r--r-- | astroid/tests/unittest_nodes.py | 39 |
4 files changed, 61 insertions, 0 deletions
@@ -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( """ |