Suggested type improvements

pyanalyze/name_check_visitor.py

@@ -109,7 +109,12 @@
     ARGS,
     KWARGS,
 )
-from .suggested_type import CallArgs, display_suggested_type
+from .suggested_type import (
+    CallArgs,
+    display_suggested_type,
+    prepare_type,
+    should_suggest_type,
+)
 from .asynq_checker import AsyncFunctionKind, AsynqChecker, FunctionInfo
 from .yield_checker import YieldChecker
 from .type_object import TypeObject, get_mro
@@ -1394,11 +1399,13 @@ def visit_FunctionDef(
                 for _, decorator in info.decorators
             )
         ):
-            self._show_error_if_checking(
-                node,
-                error_code=ErrorCode.suggested_return_type,
-                detail=display_suggested_type(return_value),
-            )
+            prepared = prepare_type(return_value)
+            if should_suggest_type(prepared):
+                self._show_error_if_checking(
+                    node,
+                    error_code=ErrorCode.suggested_return_type,
+                    detail=display_suggested_type(prepared),
+                )
 
         if evaled_function:
             return evaled_function

pyanalyze/suggested_type.py

@@ -6,7 +6,7 @@
 import ast
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Dict, Iterator, List, Mapping, Sequence, Union
+from typing import Dict, Iterator, List, Mapping, Sequence, Tuple, Union
 
 from pyanalyze.safe import safe_isinstance
 
@@ -58,12 +58,14 @@ def check(self) -> Iterator[Failure]:
             all_values = [v for v in all_values if not isinstance(v, AnyValue)]
             if not all_values:
                 continue
-            suggested = display_suggested_type(unite_values(*all_values))
+            suggested = unite_values(*all_values)
+            if not should_suggest_type(suggested):
+                continue
             failure = self.ctx.show_error(
                 param,
                 f"Suggested type for parameter {param.arg}",
                 ErrorCode.suggested_parameter_type,
-                detail=suggested,
+                detail=display_suggested_type(suggested),
                 # Otherwise we record it twice in tests. We should ultimately
                 # refactor error tracking to make it less hacky for things that
                 # show errors outside of files.
@@ -109,6 +111,15 @@ def display_suggested_type(value: Value) -> str:
     return str(cae)
 
 
+def should_suggest_type(value: Value) -> bool:
+    if isinstance(value, AnyValue):
+        return False
+    if isinstance(value, MultiValuedValue) and len(value.vals) > 5:
+        # Big unions probably aren't useful
+        return False
+    return True
+
+
 def prepare_type(value: Value) -> Value:
     """Simplify a type to turn it into a suggestion."""
     if isinstance(value, AnnotatedValue):
@@ -128,8 +139,10 @@ def prepare_type(value: Value) -> Value:
     elif isinstance(value, VariableNameValue):
         return AnyValue(AnySource.unannotated)
     elif isinstance(value, KnownValue):
-        if value.val is None or safe_isinstance(value.val, type):
+        if value.val is None:
             return value
+        elif safe_isinstance(value.val, type):
+            return SubclassValue(TypedValue(value.val))
         elif callable(value.val):
             return value  # TODO get the signature instead and return a CallableValue?
         value = replace_known_sequence_value(value)
@@ -139,22 +152,25 @@ def prepare_type(value: Value) -> Value:
             return prepare_type(value)
     elif isinstance(value, MultiValuedValue):
         vals = [prepare_type(subval) for subval in value.vals]
-        type_literals = [
-            v
-            for v in vals
-            if isinstance(v, KnownValue) and safe_isinstance(v.val, type)
-        ]
-        if len(type_literals) > 1:
-            types = [v.val for v in type_literals if isinstance(v.val, type)]
-            shared_type = get_shared_type(types)
-            type_val = SubclassValue(TypedValue(shared_type))
-            others = [
-                v
-                for v in vals
-                if not isinstance(v, KnownValue) or not safe_isinstance(v.val, type)
-            ]
-            return unite_values(type_val, *others)
-        return unite_values(*vals)
+        type_literals: List[Tuple[Value, type]] = []
+        rest: List[Value] = []
+        for subval in vals:
+            if (
+                isinstance(subval, SubclassValue)
+                and isinstance(subval.typ, TypedValue)
+                and safe_isinstance(subval.typ.typ, type)
+            ):
+                type_literals.append((subval, subval.typ.typ))
+            else:
+                rest.append(subval)
+        if type_literals:
+            shared_type = get_shared_type([typ for _, typ in type_literals])
+            if shared_type is object:
+                type_val = TypedValue(type)
+            else:
+                type_val = SubclassValue(TypedValue(shared_type))
+            return unite_values(type_val, *rest)
+        return unite_values(*[v for v, _ in type_literals], *rest)
     else:
         return value
 

pyanalyze/test_suggested_type.py

@@ -45,7 +45,6 @@ class A:
 
 
 class B(A):
-
     pass
 
 
@@ -54,7 +53,17 @@ class C(A):
 
 
 def test_prepare_type() -> None:
-    assert prepare_type(KnownValue(int) | KnownValue(str)) == SubclassValue(
-        TypedValue(object)
-    )
+    assert prepare_type(KnownValue(int) | KnownValue(str)) == TypedValue(type)
     assert prepare_type(KnownValue(C) | KnownValue(B)) == SubclassValue(TypedValue(A))
+    assert prepare_type(KnownValue(int)) == SubclassValue(TypedValue(int))
+
+    assert prepare_type(SubclassValue(TypedValue(B)) | KnownValue(C)) == SubclassValue(
+        TypedValue(A)
+    )
+    assert prepare_type(SubclassValue(TypedValue(B)) | KnownValue(B)) == SubclassValue(
+        TypedValue(B)
+    )
+    assert prepare_type(KnownValue(None) | TypedValue(str)) == KnownValue(
+        None
+    ) | TypedValue(str)
+    assert prepare_type(KnownValue(True) | KnownValue(False)) == TypedValue(bool)