From b7f96ca5a16aad5250449da54a6b235d0d5d6564 Mon Sep 17 00:00:00 2001 From: "Michael J. Sullivan" Date: Tue, 29 Oct 2019 14:43:46 -0700 Subject: [PATCH] Fix annotation of decorated methods, nested methods, nested classes (#91) The determination of the name of a function is currently done by looking at the function node and its grandparent node (for the class), and nothing else. This misses the class name for decorated methods as well as producing nonsensible results for nested classes and functions (though none of the frontends handle those right yet either). Fix this by doing a full traversal up the tree and including all functions and classes. The broken decorator behavior was also sort of accidentally causing the class name to be left off for staticmethods and classmethods (which the collector does also), so we need to keep supporting that behavior. --- .gitignore | 1 + pyannotate_tools/fixes/fix_annotate_json.py | 43 +++-- .../fixes/tests/test_annotate_json_py2.py | 150 +++++++++++++++++- .../fixes/tests/test_annotate_json_py3.py | 26 +++ 4 files changed, 204 insertions(+), 16 deletions(-) diff --git a/.gitignore b/.gitignore index b6b7ec2..6489888 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,4 @@ __pycache__ /env* *~ +.mypy_cache/ diff --git a/pyannotate_tools/fixes/fix_annotate_json.py b/pyannotate_tools/fixes/fix_annotate_json.py index 13ec449..409ed24 100644 --- a/pyannotate_tools/fixes/fix_annotate_json.py +++ b/pyannotate_tools/fixes/fix_annotate_json.py @@ -82,22 +82,26 @@ def get_init_file(dir): return f return None -def get_funcname(name, node): - # type: (Leaf, Node) -> Text - """Get function name by the following rules: +def get_funcname(node): + # type: (Optional[Node]) -> Text + """Get function name by (approximately) the following rules: - function -> function_name - - instance method -> ClassName.function_name + - method -> ClassName.function_name + + More specifically, we include every class and function name that + the node is a child of, so nested classes and functions get names like + OuterClass.InnerClass.outer_fn.inner_fn. """ - funcname = name.value - if node.parent and node.parent.parent: - grand = node.parent.parent - if grand.type == syms.classdef: - grandname = grand.children[1] - assert grandname.type == token.NAME, repr(name) - assert isinstance(grandname, Leaf) # Same as previous, for mypy - funcname = grandname.value + '.' + funcname - return funcname + components = [] # type: List[str] + while node: + if node.type in (syms.classdef, syms.funcdef): + name = node.children[1] + assert name.type == token.NAME, repr(name) + assert isinstance(name, Leaf) # Same as previous, for mypy + components.append(name.value) + node = node.parent + return '.'.join(reversed(components)) def count_args(node, results): # type: (Node, Dict[str, Base]) -> Tuple[int, bool, bool, bool] @@ -172,8 +176,19 @@ def make_annotation(self, node, results): name = results['name'] assert isinstance(name, Leaf), repr(name) assert name.type == token.NAME, repr(name) - funcname = get_funcname(name, node) + funcname = get_funcname(node) res = self.get_annotation_from_stub(node, results, funcname) + + # If we couldn't find an annotation and this is a classmethod or + # staticmethod, try again with just the funcname, since the + # type collector can't figure out class names for those. + # (We try with the full name above first so that tools that *can* figure + # that out, like dmypy suggest, can use it.) + if not res: + decs = self.get_decorators(node) + if 'staticmethod' in decs or 'classmethod' in decs: + res = self.get_annotation_from_stub(node, results, name.value) + return res stub_json_file = os.getenv('TYPE_COLLECTION_JSON') diff --git a/pyannotate_tools/fixes/tests/test_annotate_json_py2.py b/pyannotate_tools/fixes/tests/test_annotate_json_py2.py index a4c8d7a..9e59335 100644 --- a/pyannotate_tools/fixes/tests/test_annotate_json_py2.py +++ b/pyannotate_tools/fixes/tests/test_annotate_json_py2.py @@ -60,6 +60,102 @@ def nop(foo, bar): """ self.check(a, b) + def test_decorator_func(self): + self.setTestData( + [{"func_name": "foo", + "path": "", + "line": 2, + "signature": { + "arg_types": [], + "return_type": "int"}, + }]) + a = """\ + @dec + def foo(): + return 42 + """ + b = """\ + @dec + def foo(): + # type: () -> int + return 42 + """ + self.check(a, b) + + def test_decorator_method(self): + self.setTestData( + [{"func_name": "Bar.foo", + "path": "", + "line": 3, + "signature": { + "arg_types": [], + "return_type": "int"}, + }]) + a = """\ + class Bar: + @dec + @dec2 + def foo(self): + return 42 + """ + b = """\ + class Bar: + @dec + @dec2 + def foo(self): + # type: () -> int + return 42 + """ + self.check(a, b) + + def test_nested_class_func(self): + self.setTestData( + [{"func_name": "A.B.foo", + "path": "", + "line": 3, + "signature": { + "arg_types": ['str'], + "return_type": "int"}, + }]) + a = """\ + class A: + class B: + def foo(x): + return 42 + """ + b = """\ + class A: + class B: + def foo(x): + # type: (str) -> int + return 42 + """ + self.check(a, b) + + def test_nested_func(self): + self.setTestData( + [{"func_name": "A.foo.bar", + "path": "", + "line": 3, + "signature": { + "arg_types": [], + "return_type": "int"}, + }]) + a = """\ + class A: + def foo(): + def bar(): + return 42 + """ + b = """\ + class A: + def foo(): + def bar(): + # type: () -> int + return 42 + """ + self.check(a, b) + def test_keyword_only_argument(self): self.setTestData( [{"func_name": "nop", @@ -432,7 +528,7 @@ def yep(a): self.check(a, b) def test_classmethod(self): - # Class method names currently are returned without class name + # Class methods need to work without a class name self.setTestData( [{"func_name": "nop", "path": "", @@ -456,8 +552,33 @@ def nop(cls, a): """ self.check(a, b) + def test_classmethod_named(self): + # Class methods also should work *with* a class name + self.setTestData( + [{"func_name": "C.nop", + "path": "", + "line": 3, + "signature": { + "arg_types": ["int"], + "return_type": "int"} + }]) + a = """\ + class C: + @classmethod + def nop(cls, a): + return a + """ + b = """\ + class C: + @classmethod + def nop(cls, a): + # type: (int) -> int + return a + """ + self.check(a, b) + def test_staticmethod(self): - # Static method names currently are returned without class name + # Static methods need to work without a class name self.setTestData( [{"func_name": "nop", "path": "", @@ -481,6 +602,31 @@ def nop(a): """ self.check(a, b) + def test_staticmethod_named(self): + # Static methods also should work *with* a class name + self.setTestData( + [{"func_name": "C.nop", + "path": "", + "line": 3, + "signature": { + "arg_types": ["int"], + "return_type": "int"} + }]) + a = """\ + class C: + @staticmethod + def nop(a): + return a + """ + b = """\ + class C: + @staticmethod + def nop(a): + # type: (int) -> int + return a + """ + self.check(a, b) + def test_long_form(self): self.maxDiff = None self.setTestData( diff --git a/pyannotate_tools/fixes/tests/test_annotate_json_py3.py b/pyannotate_tools/fixes/tests/test_annotate_json_py3.py index 0e4d238..87f99c0 100644 --- a/pyannotate_tools/fixes/tests/test_annotate_json_py3.py +++ b/pyannotate_tools/fixes/tests/test_annotate_json_py3.py @@ -4,6 +4,8 @@ import json import os import tempfile +import unittest +import sys from lib2to3.tests.test_fixers import FixerTestCase @@ -606,3 +608,27 @@ def nop(a): return 0 def nop(a: Tuple[int, ...]) -> int: return 0 """ self.check(a, b) + + @unittest.skipIf(sys.version_info < (3, 5), 'async not supported on old python') + def test_nested_class_async_func(self): + self.setTestData( + [{"func_name": "A.B.foo", + "path": "", + "line": 3, + "signature": { + "arg_types": ['str'], + "return_type": "int"}, + }]) + a = """\ + class A: + class B: + async def foo(x): + return 42 + """ + b = """\ + class A: + class B: + async def foo(x: str) -> int: + return 42 + """ + self.check(a, b)