Skip to content

Commit 3420ef1

Browse files
authored
Fix ternary union for literals (#18023)
Fixes #18021 When I switched to (almost) always using unions as inferred from ternary expressions, I had a choice, because before we used either full context (i.e. l.h.s.) or the if type to infer the else type. After some playing I found the second one usually works better. But as we see this is not always the case, so I add some special-casing.
1 parent eb0575e commit 3420ef1

File tree

2 files changed

+30
-0
lines changed

2 files changed

+30
-0
lines changed

mypy/checkexpr.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5805,6 +5805,12 @@ def visit_conditional_expr(self, e: ConditionalExpr, allow_none_return: bool = F
58055805
context=if_type_fallback,
58065806
allow_none_return=allow_none_return,
58075807
)
5808+
5809+
# In most cases using if_type as a context for right branch gives better inferred types.
5810+
# This is however not the case for literal types, so use the full context instead.
5811+
if is_literal_type_like(full_context_else_type) and not is_literal_type_like(else_type):
5812+
else_type = full_context_else_type
5813+
58085814
res: Type = make_simplified_union([if_type, else_type])
58095815
if has_uninhabited_component(res) and not isinstance(
58105816
get_proper_type(self.type_context[-1]), UnionType

test-data/unit/check-literal.test

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2960,3 +2960,27 @@ class C(B[Literal["word"]]):
29602960
reveal_type(C().collection) # N: Revealed type is "builtins.list[Literal['word']]"
29612961
reveal_type(C().word) # N: Revealed type is "Literal['word']"
29622962
[builtins fixtures/tuple.pyi]
2963+
2964+
[case testLiteralTernaryUnionNarrowing]
2965+
from typing_extensions import Literal
2966+
from typing import Optional
2967+
2968+
SEP = Literal["a", "b"]
2969+
2970+
class Base:
2971+
def feed_data(
2972+
self,
2973+
sep: SEP,
2974+
) -> int:
2975+
return 0
2976+
2977+
class C(Base):
2978+
def feed_data(
2979+
self,
2980+
sep: Optional[SEP] = None,
2981+
) -> int:
2982+
if sep is None:
2983+
sep = "a" if int() else "b"
2984+
reveal_type(sep) # N: Revealed type is "Union[Literal['a'], Literal['b']]"
2985+
return super().feed_data(sep)
2986+
[builtins fixtures/tuple.pyi]

0 commit comments

Comments
 (0)