Format suggestions using the "signature" pyannotate format instead (#7415)

This bypasses pyannotate's parsing and processing of annotations
and allows us to just specify them directly. This is useful because
pyannotate doesn't support Callable.

Author: msullivan

mypy/suggestions.py

@@ -25,13 +25,14 @@
 from typing import (
     List, Optional, Tuple, Dict, Callable, Union, NamedTuple, TypeVar, Iterator,
 )
+from typing_extensions import TypedDict
 
 from mypy.state import strict_optional_set
 from mypy.types import (
     Type, AnyType, TypeOfAny, CallableType, UnionType, NoneType, Instance, TupleType,
     TypeVarType, FunctionLike,
     TypeStrVisitor, TypeTranslator,
-    is_optional, ProperType, get_proper_type,
+    is_optional, remove_optional, ProperType, get_proper_type,
 )
 from mypy.build import State, Graph
 from mypy.nodes import (
@@ -56,6 +57,11 @@
 import json
 import os
 
+
+PyAnnotateSignature = TypedDict('PyAnnotateSignature',
+                                {'return_type': str, 'arg_types': List[str]})
+
+
 Callsite = NamedTuple(
     'Callsite',
     [('path', str),
@@ -175,7 +181,7 @@ def suggest(self, function: str) -> str:
         if self.give_json:
             return self.json_suggestion(function, suggestion)
         else:
-            return suggestion
+            return self.format_signature(suggestion)
 
     def suggest_callsites(self, function: str) -> str:
         """Find a list of call sites of function."""
@@ -247,7 +253,11 @@ def get_args(self, is_method: bool,
             if default:
                 all_arg_types.append(default)
 
-            if all_arg_types:
+            if len(all_arg_types) == 1 and isinstance(get_proper_type(all_arg_types[0]), NoneType):
+                types.append(
+                    [UnionType.make_union([all_arg_types[0],
+                                           AnyType(TypeOfAny.explicit)])])
+            elif all_arg_types:
                 types.append(generate_type_combinations(all_arg_types))
             else:
                 # If we don't have anything, we'll try Any and object
@@ -311,7 +321,7 @@ def find_best(self, func: FuncDef, guesses: List[CallableType]) -> Tuple[Callabl
                    key=lambda s: (count_errors(errors[s]), self.score_callable(s)))
         return best, count_errors(errors[best])
 
-    def get_suggestion(self, function: str) -> str:
+    def get_suggestion(self, function: str) -> PyAnnotateSignature:
         """Compute a suggestion for a function.
 
         Return the type and whether the first argument should be ignored.
@@ -355,7 +365,7 @@ def get_suggestion(self, function: str) -> str:
         if self.no_errors and errors:
             raise SuggestionFailure("No annotation without errors")
 
-        return self.format_callable(mod, is_method, best)
+        return self.pyannotate_signature(mod, is_method, best)
 
     def format_args(self,
                     arg_kinds: List[List[int]],
@@ -483,7 +493,7 @@ def ensure_loaded(self, state: State) -> MypyFile:
     def builtin_type(self, s: str) -> Instance:
         return self.manager.semantic_analyzer.builtin_type(s)
 
-    def json_suggestion(self, function: str, suggestion: str) -> str:
+    def json_suggestion(self, function: str, suggestion: PyAnnotateSignature) -> 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
@@ -496,22 +506,33 @@ def json_suggestion(self, function: str, suggestion: str) -> str:
         path = os.path.abspath(self.graph[mod].xpath)
 
         obj = {
-            'type_comments': [suggestion],
+            'signature': suggestion,
             'line': node.line,
             'path': path,
             'func_name': func_name,
             'samples': 0
         }
         return json.dumps([obj], sort_keys=True)
 
-    def format_callable(self,
-                        cur_module: Optional[str], is_method: bool, typ: CallableType) -> str:
-        """Format a callable type in a way suitable as an annotation... kind of"""
+    def pyannotate_signature(
+        self,
+        cur_module: Optional[str],
+        is_method: bool,
+        typ: CallableType
+    ) -> PyAnnotateSignature:
+        """Format a callable type as a pyannotate dict"""
         start = int(is_method)
-        s = "({}) -> {}".format(
-            ", ".join([self.format_type(cur_module, t) for t in typ.arg_types[start:]]),
-            self.format_type(cur_module, typ.ret_type))
-        return s
+        return {
+            'arg_types': [self.format_type(cur_module, t) for t in typ.arg_types[start:]],
+            'return_type': self.format_type(cur_module, typ.ret_type),
+        }
+
+    def format_signature(self, sig: PyAnnotateSignature) -> str:
+        """Format a callable type in a way suitable as an annotation... kind of"""
+        return "({}) -> {}".format(
+            ", ".join(sig['arg_types']),
+            sig['return_type']
+        )
 
     def format_type(self, cur_module: Optional[str], typ: Type) -> str:
         return typ.accept(TypeFormatter(cur_module, self.graph))
@@ -618,7 +639,9 @@ def visit_instance(self, t: Instance) -> str:
         elif t.args != []:
             obj += '[{}]'.format(self.list_str(t.args))
 
-        if mod == 'builtins':
+        if mod_obj == ('builtins', 'unicode'):
+            return 'Text'
+        elif mod == 'builtins':
             return obj
         else:
             delim = '.' if '.' not in obj else ':'
@@ -632,6 +655,12 @@ def visit_tuple_type(self, t: TupleType) -> str:
         s = self.list_str(t.items)
         return 'Tuple[{}]'.format(s)
 
+    def visit_union_type(self, t: UnionType) -> str:
+        if len(t.items) == 2 and is_optional(t):
+            return "Optional[{}]".format(remove_optional(t).accept(self))
+        else:
+            return super().visit_union_type(t)
+
     def visit_callable_type(self, t: CallableType) -> str:
         # TODO: use extended callables?
         if is_tricky_callable(t):

test-data/unit/fine-grained-suggest.test

@@ -64,7 +64,7 @@ def untyped(x) -> None:
     foo(x)
 [builtins fixtures/isinstancelist.pyi]
 [out]
-(Union[str, int, None], Union[int, None]) -> None
+(Union[str, int, None], Optional[int]) -> None
 ==
 
 [case testSuggestInferFunc2]
@@ -100,7 +100,7 @@ from foo import bar
 bar(None)
 [out]
 (Any) -> Any
-(None) -> None
+(Optional[Any]) -> None
 ==
 
 [case testSuggestInferFuncAny2]
@@ -185,7 +185,7 @@ def foo(s):
 from foo import foo
 foo('lol')
 [out]
-(unicode) -> unicode
+(Text) -> Text
 ==
 
 [case testSuggestInferMethod1]
@@ -211,7 +211,7 @@ def bar() -> None:
     x.foo(None)
 [builtins fixtures/isinstancelist.pyi]
 [out]
-(Union[str, int, None], Union[int, None]) -> object
+(Union[str, int, None], Optional[int]) -> object
 ==
 
 [case testSuggestInferMethod2]
@@ -238,7 +238,7 @@ def bar() -> None:
     a: Union[str, int] = x.foo(None)
 [builtins fixtures/isinstancelist.pyi]
 [out]
-(Union[str, int, None], Union[int, None]) -> Union[int, str]
+(Union[str, int, None], Optional[int]) -> Union[int, str]
 ==
 
 [case testSuggestInferMethod3]
@@ -259,7 +259,7 @@ def bar() -> None:
 [builtins fixtures/isinstancelist.pyi]
 [out]
 ==
-(Union[str, None]) -> None
+(Optional[str]) -> None
 
 [case testSuggestInferFunctionUnreachable]
 # flags: --strict-optional
@@ -304,7 +304,7 @@ def bar() -> None:
 [builtins fixtures/isinstancelist.pyi]
 [out]
 ==
-(Union[str, int, None], Union[int, None]) -> Union[int, str]
+(Union[str, int, None], Optional[int]) -> Union[int, str]
 
 [case testSuggestInferNestedMethod]
 # flags: --strict-optional
@@ -405,8 +405,8 @@ def baz() -> None:
     foo(None)
 [builtins fixtures/isinstancelist.pyi]
 [out]
-(Union[str, int, None], Union[int, None]) -> object
-(Union[str, int, None], Union[int, None]) -> None
+(Union[str, int, None], Optional[int]) -> object
+(Union[str, int, None], Optional[int]) -> None
 ==
 
 [case testSuggestInferFuncDecorator1]