From e637cc58c3751405971e46d7f54347e6f89f1e41 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edgar=20Ram=C3=ADrez=20Mondrag=C3=B3n?= Date: Mon, 2 Nov 2026 11:21:18 -0600 Subject: [PATCH 1/3] Fall back to argument-only inference when literal context is too narrow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When the return type is a bare type variable (or an optional one) and the type context is a union mixing literals with other types, the context was used to solve the type variable before looking at the arguments. This could produce a spurious error if an argument was wider than the context: ```python def none_if_one[T](value: T | Literal[1]) -> T | None: ... x: str | Literal[1] | None = 1 x = none_if_one(x) # error: expected "str | Literal[1]" ``` Now, if using the literal context produces errors, retry inference from the arguments only, and use that if it type checks and the result is compatible with the context. Fixes #22334 Signed-off-by: Edgar Ramírez Mondragón --- mypy/checkexpr.py | 48 ++++++++++++++++++++++++++--- test-data/unit/check-enum.test | 19 ++++++++++++ test-data/unit/check-literal.test | 15 +++++++++ test-data/unit/check-sentinels.test | 17 ++++++++++ 4 files changed, 95 insertions(+), 4 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 2ed757eeeaa83..12bb17aea102d 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, + speculative: bool = False, ) -> tuple[Type, Type]: """Type check a call that targets a callable value. @@ -1792,8 +1793,42 @@ 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 speculative: + callee_with_ctx = self.infer_function_type_arguments_using_context(callee, context) + ctx = self.type_context[-1] + if ctx is not None and is_literal_type_like(ctx): + # 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. 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_no_ctx = self.infer_function_type_arguments_using_context( + callee, context, allow_literal_context=False + ) + + def check_speculatively(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, + speculative=True, + ) + return ret, w.has_new_errors() + + if callee_with_ctx != callee_no_ctx and check_speculatively(callee_with_ctx)[1]: + ret, has_errors = check_speculatively(callee_no_ctx) + if not has_errors and is_subtype(ret, ctx): + callee_with_ctx = callee_no_ctx + callee = callee_with_ctx formal_to_actual = map_actuals_to_formals( arg_kinds, @@ -2112,13 +2147,16 @@ def infer_arg_types_in_context( return res 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 +2216,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..66097b432d372 100644 --- a/test-data/unit/check-sentinels.test +++ b/test-data/unit/check-sentinels.test @@ -266,3 +266,20 @@ 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] +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] From 30cb49ed0b259e372b1cc808695eee3ac5c9beeb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edgar=20Ram=C3=ADrez=20Mondrag=C3=B3n?= Date: Mon, 2 Nov 2026 11:52:17 -0600 Subject: [PATCH 2/3] Add issue link to sentinel test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Edgar Ramírez Mondragón --- test-data/unit/check-sentinels.test | 1 + 1 file changed, 1 insertion(+) diff --git a/test-data/unit/check-sentinels.test b/test-data/unit/check-sentinels.test index 66097b432d372..3e85c63ba6804 100644 --- a/test-data/unit/check-sentinels.test +++ b/test-data/unit/check-sentinels.test @@ -268,6 +268,7 @@ 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 From af190c1fa717482305b0decafb0a0b22a11afc6e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edgar=20Ram=C3=ADrez=20Mondrag=C3=B3n?= Date: Mon, 2 Nov 2026 13:25:14 -0600 Subject: [PATCH 3/3] Refactor logic to its own method MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Edgar Ramírez Mondragón --- mypy/checkexpr.py | 103 +++++++++++++++++++++++++++++----------------- 1 file changed, 66 insertions(+), 37 deletions(-) diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 12bb17aea102d..f758b4434586f 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py @@ -1703,7 +1703,7 @@ def check_callable_call( callable_node: Expression | None, callable_name: str | None, object_type: Type | None, - speculative: bool = False, + skip_context_inference: bool = False, ) -> tuple[Type, Type]: """Type check a call that targets a callable value. @@ -1793,42 +1793,17 @@ def check_callable_call( callee = freshen_function_type_vars(callee) self.freshen_cache[(context, original_callee)] = callee - if callee.is_generic() and not speculative: - callee_with_ctx = self.infer_function_type_arguments_using_context(callee, context) - ctx = self.type_context[-1] - if ctx is not None and is_literal_type_like(ctx): - # 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. 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_no_ctx = self.infer_function_type_arguments_using_context( - callee, context, allow_literal_context=False - ) - - def check_speculatively(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, - speculative=True, - ) - return ret, w.has_new_errors() - - if callee_with_ctx != callee_no_ctx and check_speculatively(callee_with_ctx)[1]: - ret, has_errors = check_speculatively(callee_no_ctx) - if not has_errors and is_subtype(ret, ctx): - callee_with_ctx = callee_no_ctx - callee = callee_with_ctx + 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, @@ -2146,6 +2121,60 @@ 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, allow_literal_context: bool = True ) -> CallableType: