Support suggesting types for class and static methods

msullivan · msullivan · commit 3153ae62dd36 · 2019-08-28T13:57:39.000-07:00
diff --git a/mypy/suggestions.py b/mypy/suggestions.py
@@ -318,8 +318,7 @@ def get_suggestion(self, function: str) -> str:
         if self.no_errors and orig_errors:
             raise SuggestionFailure("Function does not typecheck.")
 
-        # FIXME: what about static and class methods?
-        is_method = bool(node.info)
+        is_method = bool(node.info) and not node.is_static
 
         if len(node.arg_names) >= 10:
             raise SuggestionFailure("Too many arguments")
@@ -482,6 +481,9 @@ def builtin_type(self, s: str) -> Instance:
     def json_suggestion(self, function: str, suggestion: str) -> str:
         """Produce a json blob for a suggestion suitable for application by pyannotate."""
         mod, func_name, node = self.find_node(function)
+        # pyannotate irritatingly drops class names for class and static methods
+        if node.is_class or node.is_static:
+            func_name = func_name.split('.', 1)[-1]
 
         # pyannotate works with either paths relative to where the
         # module is rooted or with absolute paths. We produce absolute
diff --git a/test-data/unit/fine-grained-suggest.test b/test-data/unit/fine-grained-suggest.test
@@ -456,3 +456,31 @@ def bar() -> None:
 [out]
 (str) -> str
 ==
+
+[case testSuggestInferClassMethod]
+# flags: --strict-optional
+# suggest: foo.F.bar
+# suggest: foo.F.baz
+[file foo.py]
+class F:
+    @classmethod
+    def bar(cls, x, y):
+        return x
+
+    @staticmethod
+    def baz(x, y):
+        return x
+
+[file bar.py]
+from typing import Any
+any: Any
+
+from foo import F
+def bar() -> None:
+    F.bar(0, any)
+    F.baz(0, any)
+[builtins fixtures/classmethod.pyi]
+[out]
+(int, Any) -> int
+(int, Any) -> int
+==
