Skip to content

Commit

Permalink
Fix ternary union for literals (#18023)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
ilevkivskyi authored Oct 24, 2024
1 parent eb0575e commit 3420ef1
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 0 deletions.
6 changes: 6 additions & 0 deletions mypy/checkexpr.py
Original file line number Diff line number Diff line change
Expand Up @@ -5805,6 +5805,12 @@ def visit_conditional_expr(self, e: ConditionalExpr, allow_none_return: bool = F
context=if_type_fallback,
allow_none_return=allow_none_return,
)

# In most cases using if_type as a context for right branch gives better inferred types.
# This is however not the case for literal types, so use the full context instead.
if is_literal_type_like(full_context_else_type) and not is_literal_type_like(else_type):
else_type = full_context_else_type

res: Type = make_simplified_union([if_type, else_type])
if has_uninhabited_component(res) and not isinstance(
get_proper_type(self.type_context[-1]), UnionType
Expand Down
24 changes: 24 additions & 0 deletions test-data/unit/check-literal.test
Original file line number Diff line number Diff line change
Expand Up @@ -2960,3 +2960,27 @@ class C(B[Literal["word"]]):
reveal_type(C().collection) # N: Revealed type is "builtins.list[Literal['word']]"
reveal_type(C().word) # N: Revealed type is "Literal['word']"
[builtins fixtures/tuple.pyi]

[case testLiteralTernaryUnionNarrowing]
from typing_extensions import Literal
from typing import Optional

SEP = Literal["a", "b"]

class Base:
def feed_data(
self,
sep: SEP,
) -> int:
return 0

class C(Base):
def feed_data(
self,
sep: Optional[SEP] = None,
) -> int:
if sep is None:
sep = "a" if int() else "b"
reveal_type(sep) # N: Revealed type is "Union[Literal['a'], Literal['b']]"
return super().feed_data(sep)
[builtins fixtures/tuple.pyi]

0 comments on commit 3420ef1

Please sign in to comment.