Extend special case for context-based typevar inference to typevar unions in return position (#18976)

sterliakov · web-flow · commit a0307b53ec18 · 2025-05-15T02:03:30.000-07:00
* Fixes #17221. * Fixes #17654. * Fixes #17553. * Fixes #17536. * Fixes #16659. * Fixes #16267. * Fixes #15755. * Fixes #15150. * Fixes #14664. * Incidentally improves error message in #12156. * Fixes #12092. * Fixes #11985. * Improves #11455 (but the problem with union `TypeVar | SomeFixedType` reported in comments there remains). * Fixes #10426. When using context, we can perform some overly optimistic inference when return type is `T1 | T2`. This breaks in important case of `builtins.min` when `default` and `key` are passed, essentially making them always incompatible. This is not the most principled approach, but let's see the primer results. This resolves quite a few issues (some of them duplicates, but some - substantially different), `min` problem was a very popular one... Diff run: sterliakov/mypy-issues#30
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
@@ -2013,7 +2013,12 @@ 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):
+        proper_ret = get_proper_type(ret_type)
+        if (
+            isinstance(proper_ret, TypeVarType)
+            or isinstance(proper_ret, UnionType)
+            and all(isinstance(get_proper_type(u), TypeVarType) for u in proper_ret.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.
diff --git a/test-data/unit/check-inference-context.test b/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]
