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)