diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 2ed757eeeaa83..f758b4434586f 100644 --- a/mypy/checkexpr.py +++ b/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 diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test index 4e8bf0a047502..511e96d17988c 100644 --- a/test-data/unit/check-enum.test +++ b/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] diff --git a/test-data/unit/check-literal.test b/test-data/unit/check-literal.test index b54db1f33fa5e..bae91967e4f76 100644 --- a/test-data/unit/check-literal.test +++ b/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] diff --git a/test-data/unit/check-sentinels.test b/test-data/unit/check-sentinels.test index 8b3aa79fca954..3e85c63ba6804 100644 --- a/test-data/unit/check-sentinels.test +++ b/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]