diff --git a/pytype/tools/annotate_ast/annotate_ast.py b/pytype/tools/annotate_ast/annotate_ast.py index 4adee203a..b569ed1ec 100644 --- a/pytype/tools/annotate_ast/annotate_ast.py +++ b/pytype/tools/annotate_ast/annotate_ast.py @@ -54,7 +54,13 @@ class AnnotateAstVisitor(traces.MatchAstVisitor): """ def visit_Name(self, node): - """Visits a `Name` node.""" + self._maybe_annotate(node) + + def visit_Attribute(self, node): + self._maybe_annotate(node) + + def _maybe_annotate(self, node): + """Annotates a node.""" try: ops = self.match(node) except NotImplementedError: diff --git a/pytype/tools/annotate_ast/annotate_ast_test.py b/pytype/tools/annotate_ast/annotate_ast_test.py index adfe4e58f..a188d18ff 100644 --- a/pytype/tools/annotate_ast/annotate_ast_test.py +++ b/pytype/tools/annotate_ast/annotate_ast_test.py @@ -46,6 +46,8 @@ def _get_node_key(self, node): if isinstance(node, ast.Name): return base + (node.id,) + elif isinstance(node, ast.Attribute): + return base + (node.attr,) else: return base @@ -66,5 +68,22 @@ def test_annotating_name(self): } self.assert_annotations_equal(expected, module) + def test_annotating_attribute(self): + source = """ + f = Foo() + x = f.Bar().bar() + """ + + module = self.annotate(source) + + expected = { + (1, 'Name', 'f'): 'Any', + (1, 'Name', 'Foo'): 'Any', + (2, 'Name', 'x'): 'Any', + (2, 'Name', 'f'): 'Any', + (2, 'Attribute', 'Bar'): 'Any', + (2, 'Attribute', 'bar'): 'Any', + } + self.assert_annotations_equal(expected, module) test_base.main(globals(), __name__ == '__main__')