Correctly handle generics when doing overload union math

Author: Michael0x2a

mypy/checkexpr.py

@@ -659,7 +659,8 @@ def check_call(self, callee: Type, args: List[Expression],
             # and return type.
 
             targets = cast(List[CallableType], targets)
-            unioned_callable = self.union_overload_matches(targets)
+            unioned_callable = self.union_overload_matches(
+                targets, args, arg_kinds, arg_names, context)
             if unioned_callable is None:
                 # If it was not possible to actually combine together the
                 # callables in a sound way, we give up and return the original
@@ -1204,14 +1205,18 @@ def overload_call_targets(self, arg_types: List[Type], arg_kinds: List[int],
                 arg_types, arg_kinds, arg_names, m, context=context)]
             return out if len(out) >= 1 else match
 
-    def union_overload_matches(self, callables: List[CallableType]) -> Optional[CallableType]:
+    def union_overload_matches(self, callables: List[CallableType],
+                               args: List[Expression],
+                               arg_kinds: List[int],
+                               arg_names: Optional[Sequence[Optional[str]]],
+                               context: Context) -> Optional[CallableType]:
         """Accepts a list of overload signatures and attempts to combine them together into a
         new CallableType consisting of the union of all of the given arguments and return types.
 
         Returns None if it is not possible to combine the different callables together in a
         sound manner."""
-
         new_args = [[] for _ in range(len(callables[0].arg_types))]  # type: List[List[Type]]
+        new_returns = []  # type: List[Type]
 
         expected_names = callables[0].arg_names
         expected_kinds = callables[0].arg_kinds
@@ -1222,13 +1227,25 @@ def union_overload_matches(self, callables: List[CallableType]) -> Optional[Call
                 # TODO: Enhance the union overload logic to handle a wider variety of signatures.
                 return None
 
+            if target.is_generic():
+                formal_to_actual = map_actuals_to_formals(
+                    arg_kinds, arg_names,
+                    target.arg_kinds, target.arg_names,
+                    lambda i: self.accept(args[i]))
+
+                target = freshen_function_type_vars(target)
+                target = self.infer_function_type_arguments_using_context(target, context)
+                target = self.infer_function_type_arguments(
+                    target, args, arg_kinds, formal_to_actual, context)
+
             for i, arg in enumerate(target.arg_types):
                 new_args[i].append(arg)
+            new_returns.append(target.ret_type)
 
         union_count = 0
         final_args = []
-        for args in new_args:
-            new_type = UnionType.make_simplified_union(args)
+        for args_list in new_args:
+            new_type = UnionType.make_simplified_union(args_list)
             union_count += 1 if isinstance(new_type, UnionType) else 0
             final_args.append(new_type)
 
@@ -1256,7 +1273,7 @@ def union_overload_matches(self, callables: List[CallableType]) -> Optional[Call
 
         return callables[0].copy_modified(
             arg_types=final_args,
-            ret_type=UnionType.make_simplified_union([t.ret_type for t in callables]),
+            ret_type=UnionType.make_simplified_union(new_returns),
             implicit=True,
             from_overloads=True)
 

test-data/unit/check-overloading.test

@@ -1444,52 +1444,52 @@ class Child4(Parent):
 [case testOverloadWithIncompatibleMethodOverrideAndImplementation]
 from typing import overload, Union, Any
 
-class StrSub: pass
+class Sub: pass
+class A: pass
+class B: pass
 
 class ParentWithTypedImpl:
     @overload
-    def f(self, arg: int) -> int: ...
+    def f(self, arg: A) -> A: ...
     @overload
-    def f(self, arg: str) -> str: ...
-    def f(self, arg: Union[int, str]) -> Union[int, str]: ...
+    def f(self, arg: B) -> B: ...
+    def f(self, arg: Union[A, B]) -> Union[A, B]: ...
 
 class Child1(ParentWithTypedImpl):
     @overload # E: Signature of "f" incompatible with supertype "ParentWithTypedImpl"
-    def f(self, arg: int) -> int: ...
+    def f(self, arg: A) -> A: ...
     @overload
-    def f(self, arg: StrSub) -> str: ...
-    def f(self, arg: Union[int, StrSub]) -> Union[int, str]: ...
+    def f(self, arg: Sub) -> B: ...
+    def f(self, arg: Union[A, Sub]) -> Union[A, B]: ...
 
 class Child2(ParentWithTypedImpl):
     @overload # E: Signature of "f" incompatible with supertype "ParentWithTypedImpl"
-    def f(self, arg: int) -> int: ...
+    def f(self, arg: A) -> A: ...
     @overload
-    def f(self, arg: StrSub) -> str: ...
+    def f(self, arg: Sub) -> B: ...
     def f(self, arg: Any) -> Any: ...
 
 class ParentWithDynamicImpl:
     @overload
-    def f(self, arg: int) -> int: ...
+    def f(self, arg: A) -> A: ...
     @overload
-    def f(self, arg: str) -> str: ...
+    def f(self, arg: B) -> B: ...
     def f(self, arg: Any) -> Any: ...
 
 class Child3(ParentWithDynamicImpl):
     @overload # E: Signature of "f" incompatible with supertype "ParentWithDynamicImpl"
-    def f(self, arg: int) -> int: ...
+    def f(self, arg: A) -> A: ...
     @overload
-    def f(self, arg: StrSub) -> str: ...
-    def f(self, arg: Union[int, StrSub]) -> Union[int, str]: ...
+    def f(self, arg: Sub) -> B: ...
+    def f(self, arg: Union[A, Sub]) -> Union[A, B]: ...
 
 class Child4(ParentWithDynamicImpl):
     @overload # E: Signature of "f" incompatible with supertype "ParentWithDynamicImpl"
-    def f(self, arg: int) -> int: ...
+    def f(self, arg: A) -> A: ...
     @overload
-    def f(self, arg: StrSub) -> str: ...
+    def f(self, arg: Sub) -> B: ...
     def f(self, arg: Any) -> Any: ...
 
-[builtins fixtures/tuple.pyi]
-
 [case testOverloadInferUnionReturnBasic]
 from typing import overload, Union
 
@@ -1515,8 +1515,6 @@ def f2(x): ...
 
 reveal_type(f2(arg1))  # E: Revealed type is '__main__.B'
 
-[builtins fixtures/tuple.pyi]
-
 [case testOverloadInferUnionReturnMultipleArguments]
 from typing import overload, Union
 
@@ -1544,7 +1542,6 @@ def f2(x, y): ...
 reveal_type(f2(arg1, arg1))
 reveal_type(f2(arg1, C()))
 
-[builtins fixtures/tuple.pyi]
 [out]
 main:16: error: Revealed type is '__main__.B'
 main:16: error: Argument 1 to "f1" has incompatible type "Union[A, C]"; expected "A"
@@ -1553,6 +1550,24 @@ main:24: error: Revealed type is 'Union[__main__.B, __main__.D]'
 main:24: error: Argument 2 to "f2" has incompatible type "Union[A, C]"; expected "C"
 main:25: error: Revealed type is 'Union[__main__.B, __main__.D]'
 
+[case testOverloadInferUnionSkipIfParameterNamesAreDifferent]
+from typing import overload, Union
+
+class A: ...
+class B: ...
+class C: ...
+
+@overload
+def f(x: A) -> B: ...
+@overload
+def f(y: B) -> C: ...
+def f(x): ...
+
+x: Union[A, B]
+reveal_type(f(A()))  # E: Revealed type is '__main__.B'
+reveal_type(f(B()))  # E: Revealed type is '__main__.C'
+f(x)                 # E: Argument 1 to "f" has incompatible type "Union[A, B]"; expected "A"
+
 [case testOverloadInferUnionReturnFunctionsWithKwargs]
 from typing import overload, Union, Optional
 
@@ -1584,3 +1599,50 @@ main:19: error: Revealed type is '__main__.C'
 main:19: error: Argument 2 to "f" has incompatible type "Union[B, C]"; expected "Optional[B]"
 main:21: error: Revealed type is '__main__.A'
 
+[case testOverloadingInferUnionReturnWithTypevarWithValueRestriction]
+from typing import overload, Union, TypeVar, Generic
+
+class A: pass
+class B: pass
+class C: pass
+
+T = TypeVar('T', B, C)
+
+class Wrapper(Generic[T]):
+    @overload
+    def f(self, x: T) -> B: ...
+
+    @overload
+    def f(self, x: A) -> C: ...
+
+    def f(self, x): ...
+
+obj: Wrapper[B] = Wrapper()
+x: Union[A, B]
+
+reveal_type(obj.f(A()))  # E: Revealed type is '__main__.C'
+reveal_type(obj.f(B()))  # E: Revealed type is '__main__.B'
+reveal_type(obj.f(x))    # E: Revealed type is 'Union[__main__.B, __main__.C]'
+
+[case testOverloadingInferUnionReturnWithTypevarReturn]
+from typing import overload, Union, TypeVar, Generic
+
+T = TypeVar('T')
+
+class Wrapper1(Generic[T]): pass
+class Wrapper2(Generic[T]): pass
+class A: pass
+class B: pass
+
+@overload
+def f(x: Wrapper1[T]) -> T: ...
+@overload
+def f(x: Wrapper2[T]) -> T: ...
+def f(x): ...
+
+obj1: Union[Wrapper1[A], Wrapper2[A]]
+reveal_type(f(obj1))  # E: Revealed type is '__main__.A'
+
+obj2: Union[Wrapper1[A], Wrapper2[B]]
+reveal_type(f(obj2))  # E: Revealed type is 'Union[__main__.A, __main__.B]'
+