Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Narrow individual items when matching a tuple to a sequence pattern #16905

Merged
merged 3 commits into from
Apr 4, 2024
Merged

Narrow individual items when matching a tuple to a sequence pattern #16905

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
30c72e1
Narrow individual items when matching a tuple to a sequence pattern
Feb 11, 2024
1c33add
Fix annotation for Python < 3.9
Feb 11, 2024
92cb2c5
Add test cases with mismatching subject and target lengths
Apr 2, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions mypy/checker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5086,6 +5086,9 @@ def visit_match_stmt(self, s: MatchStmt) -> None:
)
self.remove_capture_conflicts(pattern_type.captures, inferred_types)
self.push_type_map(pattern_map)
if pattern_map:
for expr, typ in pattern_map.items():
self.push_type_map(self._get_recursive_sub_patterns_map(expr, typ))
self.push_type_map(pattern_type.captures)
if g is not None:
with self.binder.frame_context(can_skip=False, fall_through=3):
Expand Down Expand Up @@ -5123,6 +5126,20 @@ def visit_match_stmt(self, s: MatchStmt) -> None:
with self.binder.frame_context(can_skip=False, fall_through=2):
pass

def _get_recursive_sub_patterns_map(
self, expr: Expression, typ: Type
) -> dict[Expression, Type]:
sub_patterns_map: dict[Expression, Type] = {}
typ_ = get_proper_type(typ)
if isinstance(expr, TupleExpr) and isinstance(typ_, TupleType):
# When matching a tuple expression with a sequence pattern, narrow individual tuple items
assert len(expr.items) == len(typ_.items)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add some test cases where the length doesn't match, to make sure we don't fail this assertion? Especially in nested tuples.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done! I did not check anything in non-matching cases, just that the assertion didn't fail, eg.

match m, n, o:
    case [3, "foo"]:
        pass
    case [3, "foo", True, True]:
        pass

Interestingly, these paths are not considered unreachable by mypy. They probably should, but that's unrelated to this issue!

for item_expr, item_typ in zip(expr.items, typ_.items):
sub_patterns_map[item_expr] = item_typ
sub_patterns_map.update(self._get_recursive_sub_patterns_map(item_expr, item_typ))

return sub_patterns_map

def infer_variable_types_from_type_maps(self, type_maps: list[TypeMap]) -> dict[Var, Type]:
all_captures: dict[Var, list[tuple[NameExpr, Type]]] = defaultdict(list)
for tm in type_maps:
Expand Down
38 changes: 38 additions & 0 deletions test-data/unit/check-python310.test
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -341,6 +341,44 @@ match m:
reveal_type(m) # N: Revealed type is "builtins.list[builtins.list[builtins.str]]"
[builtins fixtures/list.pyi]

[case testMatchSequencePatternNarrowSubjectItems]
m: int
n: str
o: bool

match m, n, o:
case [3, "foo", True]:
reveal_type(m) # N: Revealed type is "Literal[3]"
reveal_type(n) # N: Revealed type is "Literal['foo']"
reveal_type(o) # N: Revealed type is "Literal[True]"
case [a, b, c]:
reveal_type(m) # N: Revealed type is "builtins.int"
reveal_type(n) # N: Revealed type is "builtins.str"
reveal_type(o) # N: Revealed type is "builtins.bool"

reveal_type(m) # N: Revealed type is "builtins.int"
reveal_type(n) # N: Revealed type is "builtins.str"
reveal_type(o) # N: Revealed type is "builtins.bool"
[builtins fixtures/tuple.pyi]

[case testMatchSequencePatternNarrowSubjectItemsRecursive]
m: int
n: int
o: int
p: int
q: int
r: int

match m, (n, o), (p, (q, r)):
case [0, [1, 2], [3, [4, 5]]]:
reveal_type(m) # N: Revealed type is "Literal[0]"
reveal_type(n) # N: Revealed type is "Literal[1]"
reveal_type(o) # N: Revealed type is "Literal[2]"
reveal_type(p) # N: Revealed type is "Literal[3]"
reveal_type(q) # N: Revealed type is "Literal[4]"
reveal_type(r) # N: Revealed type is "Literal[5]"
[builtins fixtures/tuple.pyi]

-- Mapping Pattern --

[case testMatchMappingPatternCaptures]
Expand Down
Loading