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

@@ -6,3 +6,4 @@
 __pycache__
 /env*
 *~
+.mypy_cache/

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')

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": "<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",
@@ -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>",
@@ -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>",
@@ -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(

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": "<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)