Extend special-case for context-based typevar inference in return position to typevar unions

Author: sterliakov

mypy/checkexpr.py

@@ -2013,7 +2013,11 @@ def infer_function_type_arguments_using_context(
             #     variables in an expression are inferred at the same time.
             #     (And this is hard, also we need to be careful with lambdas that require
             #     two passes.)
-        if isinstance(ret_type, TypeVarType):
+        if (
+            isinstance(ret_type, TypeVarType)
+            or isinstance(ret_type, UnionType)
+            and all(isinstance(u, TypeVarType) for u in ret_type.items)
+        ):
             # Another special case: the return type is a type variable. If it's unrestricted,
             # we could infer a too general type for the type variable if we use context,
             # and this could result in confusing and spurious type errors elsewhere.

test-data/unit/check-inference-context.test

@@ -1495,3 +1495,18 @@ def g(b: Optional[str]) -> None:
         z: Callable[[], str] = lambda: reveal_type(b)  # N: Revealed type is "builtins.str"
         f2(lambda: reveal_type(b))  # N: Revealed type is "builtins.str"
         lambda: reveal_type(b)  # N: Revealed type is "builtins.str"
+
+[case testInferenceContextReturningTypeVarUnion]
+from collections.abc import Callable, Iterable
+from typing import TypeVar, Union
+
+_T1 = TypeVar("_T1")
+_T2 = TypeVar("_T2")
+
+def mymin(
+    iterable: Iterable[_T1], /, *, key: Callable[[_T1], int], default: _T2
+) -> Union[_T1, _T2]: ...
+
+def check(paths: Iterable[str], key: Callable[[str], int]) -> Union[str, None]:
+    return mymin(paths, key=key, default=None)
+[builtins fixtures/tuple.pyi]