Fall back to argument-only inference when literal context is too narrow

mypy/checkexpr.py

@@ -1703,6 +1703,7 @@ def check_callable_call(
         callable_node: Expression | None,
         callable_name: str | None,
         object_type: Type | None,
+        skip_context_inference: bool = False,
     ) -> tuple[Type, Type]:
         """Type check a call that targets a callable value.
 
@@ -1792,8 +1793,17 @@ def check_callable_call(
                     callee = freshen_function_type_vars(callee)
                 self.freshen_cache[(context, original_callee)] = callee
 
-        if callee.is_generic():
-            callee = self.infer_function_type_arguments_using_context(callee, context)
+        if callee.is_generic() and not skip_context_inference:
+            callee = self.infer_callee_using_context(
+                callee,
+                args,
+                arg_kinds,
+                context,
+                arg_names,
+                callable_node,
+                callable_name,
+                object_type,
+            )
 
         formal_to_actual = map_actuals_to_formals(
             arg_kinds,
@@ -2111,14 +2121,71 @@ def infer_arg_types_in_context(
                 res.append(self.accept(arg))
         return res
 
+    def infer_callee_using_context(
+        self,
+        callee: CallableType,
+        args: list[Expression],
+        arg_kinds: list[ArgKind],
+        context: Context,
+        arg_names: Sequence[str | None] | None,
+        callable_node: Expression | None,
+        callable_name: str | None,
+        object_type: Type | None,
+    ) -> CallableType:
+        """Infer type variables of a generic callee using the type context.
+
+        A literal-like context is meant to propagate literal types down, but it can be
+        too narrow for a union that mixes literals with other types. In that case, fall
+        back to inferring from the arguments only if using the context results in errors,
+        and the result type checks and is compatible with the context.
+        """
+        callee_with_ctx = self.infer_function_type_arguments_using_context(callee, context)
+        ctx = self.type_context[-1]
+        if ctx is None or not is_literal_type_like(ctx):
+            return callee_with_ctx
+        callee_no_ctx = self.infer_function_type_arguments_using_context(
+            callee, context, allow_literal_context=False
+        )
+        if callee_with_ctx == callee_no_ctx:
+            return callee_with_ctx
+
+        def try_call(c: CallableType) -> tuple[Type, bool]:
+            """Check the call with callee c, discarding errors and inferred types.
+
+            Return the resulting type and whether there were any errors.
+            """
+            with self.msg.filter_errors() as w, self.chk.local_type_map:
+                ret, _ = self.check_callable_call(
+                    c,
+                    args,
+                    arg_kinds,
+                    context,
+                    arg_names,
+                    callable_node,
+                    callable_name,
+                    object_type,
+                    skip_context_inference=True,
+                )
+            return ret, w.has_new_errors()
+
+        if not try_call(callee_with_ctx)[1]:
+            return callee_with_ctx
+        ret, has_errors = try_call(callee_no_ctx)
+        if not has_errors and is_subtype(ret, ctx):
+            return callee_no_ctx
+        return callee_with_ctx
+
     def infer_function_type_arguments_using_context(
-        self, callable: CallableType, error_context: Context
+        self, callable: CallableType, error_context: Context, allow_literal_context: bool = True
     ) -> CallableType:
         """Unify callable return type to type context to infer type vars.
 
         For example, if the return type is set[t] where 't' is a type variable
         of callable, and if the context is set[int], return callable modified
         by substituting 't' with 'int'.
+
+        If allow_literal_context is False, don't use a literal-like context when the
+        return type is a bare type variable (see below).
         """
         ctx = self.type_context[-1]
         if not ctx:
@@ -2178,7 +2245,9 @@ def infer_function_type_arguments_using_context(
             #        expects_literal(identity(3))  # Should type-check
             # TODO: we may want to add similar exception if all arguments are lambdas, since
             # in this case external context is almost everything we have.
-            if not is_generic_instance(ctx) and not is_literal_type_like(ctx):
+            if not is_generic_instance(ctx) and not (
+                allow_literal_context and is_literal_type_like(ctx)
+            ):
                 return callable.copy_modified()
         args = infer_type_arguments(
             callable.variables, ret_type, erased_ctx, skip_unsatisfied=True

test-data/unit/check-enum.test

@@ -3007,3 +3007,22 @@ match foo:
         foo()
         reveal_type(foo) # N: Revealed type is "Literal[__main__.Foo.FOO2]"
 [builtins fixtures/primitives.pyi]
+
+[case testEnumLiteralInUnionContextWithGenericFallback]
+# https://github.com/python/mypy/issues/22334
+from enum import Enum
+from typing import Literal, TypeVar
+
+T = TypeVar("T")
+
+class E(Enum):
+    SPECIAL = 1
+
+def none_if_special(value: T | Literal[E.SPECIAL]) -> T | None:
+    return None if value is E.SPECIAL else value
+
+def func() -> None:
+    x: str | Literal[E.SPECIAL] | None = E.SPECIAL
+    x = none_if_special(x)
+    reveal_type(x)  # N: Revealed type is "builtins.str | None"
+[builtins fixtures/tuple.pyi]

test-data/unit/check-literal.test

@@ -3164,3 +3164,18 @@ bar: Literal[
 ]
 reveal_type(bar)  # N: Revealed type is "Literal['duration[s][pyarrow]'] | Literal['duration[ms][pyarrow]']"
 [builtins fixtures/primitives.pyi]
+
+[case testLiteralInUnionContextWithGenericFallback]
+# https://github.com/python/mypy/issues/22334
+from typing import Literal, TypeVar
+
+T = TypeVar("T")
+
+def none_if_one(value: T | Literal[1]) -> T | None:
+    return None if isinstance(value, int) else value
+
+def func() -> None:
+    x: str | Literal[1] | None = 1
+    x = none_if_one(x)
+    reveal_type(x)  # N: Revealed type is "builtins.str | None"
+[builtins fixtures/tuple.pyi]

test-data/unit/check-sentinels.test

@@ -266,3 +266,21 @@ reveal_type(X)  # N: Revealed type is "X"
 reveal_type(inspect_sentinel)  # N: Revealed type is "def (arg: typing_extensions.Sentinel)"
 inspect_sentinel(arg=X)
 [builtins fixtures/tuple.pyi]
+
+[case testSentinelInUnionContextWithGenericFallback]
+# https://github.com/python/mypy/issues/22334
+from typing import Generic, TypeVar, reveal_type
+from typing_extensions import sentinel
+
+Undefined = sentinel("Undefined")
+
+_T = TypeVar("_T")
+
+def none_if_undefined(value: _T | Undefined) -> _T | None:
+    return None if value is Undefined else value
+
+def func() -> None:
+    device_id: str | Undefined | None = Undefined
+    device_id = none_if_undefined(device_id)
+    reveal_type(device_id)  # N: Revealed type is "builtins.str | None"
+[builtins fixtures/tuple.pyi]