Skip to content

Commit 6e21887

Browse files
authored
Narrow tagged unions in match statements (#18791)
Fixes #16286. --- This PR was generated by an AI system in collaboration with maintainers: @hauntsaninja --------- Signed-off-by: Gene Parmesan Thomas <[email protected]> Signed-off-by: gopoto <[email protected]>
1 parent e37d92d commit 6e21887

File tree

2 files changed

+50
-0
lines changed

2 files changed

+50
-0
lines changed

mypy/checker.py

+2
Original file line numberDiff line numberDiff line change
@@ -5527,6 +5527,8 @@ def visit_match_stmt(self, s: MatchStmt) -> None:
55275527
pattern_map, else_map = conditional_types_to_typemaps(
55285528
named_subject, pattern_type.type, pattern_type.rest_type
55295529
)
5530+
pattern_map = self.propagate_up_typemap_info(pattern_map)
5531+
else_map = self.propagate_up_typemap_info(else_map)
55305532
self.remove_capture_conflicts(pattern_type.captures, inferred_types)
55315533
self.push_type_map(pattern_map, from_assignment=False)
55325534
if pattern_map:

test-data/unit/check-python310.test

+48
Original file line numberDiff line numberDiff line change
@@ -332,6 +332,54 @@ match [SubClass("a"), SubClass("b")]:
332332
reveal_type(rest) # N: Revealed type is "builtins.list[__main__.Example]"
333333
[builtins fixtures/tuple.pyi]
334334

335+
# Narrowing union-based values via a literal pattern on an indexed/attribute subject
336+
# -------------------------------------------------------------------------------
337+
# Literal patterns against a union of types can be used to narrow the subject
338+
# itself, not just the expression being matched. Previously, the patterns below
339+
# failed to narrow the `d` variable, leading to errors for missing members; we
340+
# now propagate the type information up to the parent.
341+
342+
[case testMatchNarrowingUnionTypedDictViaIndex]
343+
from typing import Literal, TypedDict
344+
345+
class A(TypedDict):
346+
tag: Literal["a"]
347+
name: str
348+
349+
class B(TypedDict):
350+
tag: Literal["b"]
351+
num: int
352+
353+
d: A | B
354+
match d["tag"]:
355+
case "a":
356+
reveal_type(d) # N: Revealed type is "TypedDict('__main__.A', {'tag': Literal['a'], 'name': builtins.str})"
357+
reveal_type(d["name"]) # N: Revealed type is "builtins.str"
358+
case "b":
359+
reveal_type(d) # N: Revealed type is "TypedDict('__main__.B', {'tag': Literal['b'], 'num': builtins.int})"
360+
reveal_type(d["num"]) # N: Revealed type is "builtins.int"
361+
[typing fixtures/typing-typeddict.pyi]
362+
363+
[case testMatchNarrowingUnionClassViaAttribute]
364+
from typing import Literal
365+
366+
class A:
367+
tag: Literal["a"]
368+
name: str
369+
370+
class B:
371+
tag: Literal["b"]
372+
num: int
373+
374+
d: A | B
375+
match d.tag:
376+
case "a":
377+
reveal_type(d) # N: Revealed type is "__main__.A"
378+
reveal_type(d.name) # N: Revealed type is "builtins.str"
379+
case "b":
380+
reveal_type(d) # N: Revealed type is "__main__.B"
381+
reveal_type(d.num) # N: Revealed type is "builtins.int"
382+
335383
[case testMatchSequenceUnion-skip]
336384
from typing import List, Union
337385
m: Union[List[List[str]], str]

0 commit comments

Comments
 (0)