Skip to content

Commit

Permalink
Fix annotation of decorated methods, nested methods, nested classes (#91
Browse files Browse the repository at this point in the history
)

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.
  • Loading branch information
msullivan committed Oct 29, 2019
1 parent ce8afa9 commit b7f96ca
Show file tree
Hide file tree
Showing 4 changed files with 204 additions and 16 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -6,3 +6,4 @@
__pycache__
/env*
*~
.mypy_cache/
43 changes: 29 additions & 14 deletions pyannotate_tools/fixes/fix_annotate_json.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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')
Expand Down
150 changes: 148 additions & 2 deletions pyannotate_tools/fixes/tests/test_annotate_json_py2.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,102 @@ def nop(foo, bar):
"""
self.check(a, b)

def test_decorator_func(self):
self.setTestData(
[{"func_name": "foo",
"path": "<string>",
"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": "<string>",
"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": "<string>",
"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": "<string>",
"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",
Expand Down Expand Up @@ -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": "<string>",
Expand All @@ -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": "<string>",
"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": "<string>",
Expand All @@ -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": "<string>",
"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(
Expand Down
26 changes: 26 additions & 0 deletions pyannotate_tools/fixes/tests/test_annotate_json_py3.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
import json
import os
import tempfile
import unittest
import sys

from lib2to3.tests.test_fixers import FixerTestCase

Expand Down Expand Up @@ -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": "<string>",
"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)

0 comments on commit b7f96ca

Please sign in to comment.