Handle more subject expressions

Author: cdce8p

mypy/checkpattern.py

@@ -16,12 +16,15 @@
 from mypy.messages import MessageBuilder
 from mypy.nodes import (
     ARG_POS,
+    AssignmentExpr,
     Context,
     Expression,
     IndexExpr,
     IntExpr,
+    ListExpr,
     MemberExpr,
     NameExpr,
+    TupleExpr,
     TypeAlias,
     TypeInfo,
     UnaryExpr,
@@ -194,7 +197,8 @@ def visit_or_pattern(self, o: OrPattern) -> PatternType:
         capture_expr_keys: set[Key | None] = set()
         # Collect captures from the first subpattern
         for expr, typ in pattern_types[0].captures.items():
-            node = get_var(expr)
+            if (node := get_var(expr)) is None:
+                continue
             key = literal_hash(expr)
             capture_types[node][key].append((expr, typ))
             if isinstance(expr, NameExpr):
@@ -209,7 +213,8 @@ def visit_or_pattern(self, o: OrPattern) -> PatternType:
                 # Only fail for directly captured names (with NameExpr)
                 self.msg.fail(message_registry.OR_PATTERN_ALTERNATIVE_NAMES, o.patterns[i])
             for expr, typ in pattern_type.captures.items():
-                node = get_var(expr)
+                if (node := get_var(expr)) is None:
+                    continue
                 key = literal_hash(expr)
                 capture_types[node][key].append((expr, typ))
 
@@ -312,15 +317,26 @@ def visit_sequence_pattern(self, o: SequencePattern) -> PatternType:
             inner_types, star_position, required_patterns
         )
         current_subjects: list[list[Expression]] = [[] for _ in range(len(contracted_inner_types))]
+        end_pos = len(contracted_inner_types) if star_position is None else star_position
         for subject in self.subject_context[-1]:
-            # Support x[0], x[1], ... lookup until wildcard
-            end_pos = len(contracted_inner_types) if star_position is None else star_position
-            for i in range(end_pos):
-                current_subjects[i].append(IndexExpr(subject, IntExpr(i)))
-            # For everything after wildcard use x[-2], x[-1]
-            for i in range((star_position or -1) + 1, len(contracted_inner_types)):
-                offset = len(contracted_inner_types) - i
-                current_subjects[i].append(IndexExpr(subject, UnaryExpr("-", IntExpr(offset))))
+            if isinstance(subject, (ListExpr, TupleExpr)):
+                # For list and tuple expressions, lookup expression in items
+                for i in range(end_pos):
+                    if i < len(subject.items):
+                        current_subjects[i].append(subject.items[i])
+                if star_position is not None:
+                    for i in range(star_position + 1, len(contracted_inner_types)):
+                        offset = len(contracted_inner_types) - i
+                        if offset <= len(subject.items):
+                            current_subjects[i].append(subject.items[-offset])
+            else:
+                # Support x[0], x[1], ... lookup until wildcard
+                for i in range(end_pos):
+                    current_subjects[i].append(IndexExpr(subject, IntExpr(i)))
+                # For everything after wildcard use x[-2], x[-1]
+                for i in range((star_position or -1) + 1, len(contracted_inner_types)):
+                    offset = len(contracted_inner_types) - i
+                    current_subjects[i].append(IndexExpr(subject, UnaryExpr("-", IntExpr(offset))))
         for p, t, s in zip(o.patterns, contracted_inner_types, current_subjects):
             pattern_type = self.accept(p, t, s)
             typ, rest, type_map = pattern_type
@@ -794,7 +810,8 @@ def update_type_map(
         already_captured = {literal_hash(expr) for expr in original_type_map}
         for expr, typ in extra_type_map.items():
             if literal_hash(expr) in already_captured:
-                node = get_var(expr)
+                if (node := get_var(expr)) is None:
+                    continue
                 self.msg.fail(
                     message_registry.MULTIPLE_ASSIGNMENTS_IN_PATTERN.format(node.name), expr
                 )
@@ -849,7 +866,7 @@ def get_match_arg_names(typ: TupleType) -> list[str | None]:
     return args
 
 
-def get_var(expr: Expression) -> Var:
+def get_var(expr: Expression) -> Var | None:
     """
     Warning: this in only true for expressions captured by a match statement.
     Don't call it from anywhere else
@@ -858,7 +875,10 @@ def get_var(expr: Expression) -> Var:
         return get_var(expr.expr)
     if isinstance(expr, IndexExpr):
         return get_var(expr.base)
-    assert isinstance(expr, NameExpr), expr
+    if isinstance(expr, AssignmentExpr):
+        return get_var(expr.target)
+    if not isinstance(expr, NameExpr):
+        return None
     node = expr.node
     assert isinstance(node, Var), node
     return node

mypy/fastparse.py

@@ -398,6 +398,7 @@ def __init__(
         # 'C' for class, 'D' for function signature, 'F' for function, 'L' for lambda
         self.class_and_function_stack: list[Literal["C", "D", "F", "L"]] = []
         self.imports: list[ImportBase] = []
+        self.match_stmt_subject = False
 
         self.options = options
         self.is_stub = is_stub
@@ -1760,7 +1761,7 @@ def visit_Name(self, n: Name) -> NameExpr:
     # List(expr* elts, expr_context ctx)
     def visit_List(self, n: ast3.List) -> ListExpr | TupleExpr:
         expr_list: list[Expression] = [self.visit(e) for e in n.elts]
-        if isinstance(n.ctx, ast3.Store):
+        if isinstance(n.ctx, ast3.Store) or self.match_stmt_subject:
             # [x, y] = z and (x, y) = z means exactly the same thing
             e: ListExpr | TupleExpr = TupleExpr(expr_list)
         else:
@@ -1793,8 +1794,11 @@ def visit_Index(self, n: Index) -> Node:
 
     # Match(expr subject, match_case* cases) # python 3.10 and later
     def visit_Match(self, n: Match) -> MatchStmt:
+        self.match_stmt_subject = True
+        subject = self.visit(n.subject)
+        self.match_stmt_subject = False
         node = MatchStmt(
-            self.visit(n.subject),
+            subject,
             [self.visit(c.pattern) for c in n.cases],
             [self.visit(c.guard) for c in n.cases],
             [self.as_required_block(c.body) for c in n.cases],

test-data/unit/check-python310.test

@@ -318,6 +318,7 @@ match x:
         pass
 
 [case testMatchSequencePatternWithInvalidClassPattern]
+# flags: --warn-unreachable
 class Example:
     __match_args__ = ("value",)
     def __init__(self, value: str) -> None:
@@ -327,8 +328,10 @@ SubClass: type[Example]
 
 match [SubClass("a"), SubClass("b")]:
     case [SubClass(value), *rest]:  # E: Expected type in class pattern; found "type[__main__.Example]"
-        reveal_type(value)  # E: Cannot determine type of "value" \
-                            # N: Revealed type is "Any"
+        reveal_type(value)  # E: Statement is unreachable
+        reveal_type(rest)
+    case [Example(value), *rest]:
+        reveal_type(value)  # N: Revealed type is "builtins.str"
         reveal_type(rest)  # N: Revealed type is "builtins.list[__main__.Example]"
 [builtins fixtures/tuple.pyi]
 
@@ -3048,3 +3051,53 @@ match m:
         reveal_type(m.a)  # N: Revealed type is "Any"
         reveal_type(m.a[0])  # N: Revealed type is "__main__.A"
         reveal_type(m.a[0].a)  # N: Revealed type is "Union[Literal['Hello'], Literal['World']]"
+
+[case testMatchSubjectExpression]
+# flags: --warn-unreachable
+m: object
+n: object
+o: object
+def func(): ...
+
+match (m, n, o):
+    case [1, 2, 3] | [2, 3, 4]:
+        reveal_type(m)  # N: Revealed type is "Union[Literal[1], Literal[2]]"
+        reveal_type(n)  # N: Revealed type is "Union[Literal[2], Literal[3]]"
+        reveal_type(o)  # N: Revealed type is "Union[Literal[3], Literal[4]]"
+    case [1, 2, 3, 4] | [2, 3, 4, 5]:
+        # No match -> don't crash
+        reveal_type(m)  # E: Statement is unreachable
+    case [1, *_, 3] | [2, *_, 4]:
+        reveal_type(m)  # N: Revealed type is "Union[Literal[1], Literal[2]]"
+        reveal_type(n)  # N: Revealed type is "builtins.object"
+        reveal_type(o)  # N: Revealed type is "Union[Literal[3], Literal[4]]"
+    case [1, *_, 3, 4, 5] | [2, *_, 3, 4, 5]:
+        # No match -> don't crash
+        reveal_type(m)  # E: Statement is unreachable
+
+match [m, n, o]:
+    case [1, 2, 3] | [2, 3, 4]:
+        reveal_type(m)  # N: Revealed type is "Union[Literal[1], Literal[2]]"
+        reveal_type(n)  # N: Revealed type is "Union[Literal[2], Literal[3]]"
+        reveal_type(o)  # N: Revealed type is "Union[Literal[3], Literal[4]]"
+    case [1, 2, 3, 4] | [2, 3, 4, 5]:
+        # No match -> don't crash
+        reveal_type(m)  # E: Statement is unreachable
+    case [1, *_, 3] | [2, *_, 4]:
+        reveal_type(m)  # N: Revealed type is "Union[Literal[1], Literal[2]]"
+        reveal_type(n)  # N: Revealed type is "builtins.object"
+        reveal_type(o)  # N: Revealed type is "Union[Literal[3], Literal[4]]"
+    case [1, *_, 3, 4, 5] | [2, *_, 3, 4, 5]:
+        # No match -> don't crash
+        reveal_type(m)  # E: Statement is unreachable
+
+match a := m:
+    case [1, 2] | [3, 4]:
+        reveal_type(a)  # N: Revealed type is "typing.Sequence[builtins.int]"
+        reveal_type(a[0])  # N: Revealed type is "Union[Literal[1], Literal[3]]"
+
+match func():
+    # Don't crash for subject expressions which can't be narrowed
+    case [1, 2] | [3, 4]:
+        ...
+[builtins fixtures/tuple.pyi]