add match-case support for unions

Author: randolf-scholz

Lib/test/test_patma.py

@@ -4,7 +4,6 @@
 import dis
 import enum
 import inspect
-from re import I
 import sys
 import unittest
 from test import support
@@ -16,6 +15,13 @@ class Point:
     y: int
 
 
+@dataclasses.dataclass
+class Point3D:
+    x: int
+    y: int
+    z: int
+
+
 class TestCompiler(unittest.TestCase):
 
     def test_refleaks(self):
@@ -2891,11 +2897,81 @@ class B(A): ...
 
     def test_patma_union_type(self):
         IntOrStr = int | str
-        x = 0
-        match x:
+        w = None
+        match 0:
             case IntOrStr():
-                x = 1
-        self.assertEqual(x, 1)
+                w = 0
+        self.assertEqual(w, 0)
+
+    def test_patma_union_no_match(self):
+        StrOrBytes = str | bytes
+        w = None
+        match 0:
+            case StrOrBytes():
+                w = 0
+        self.assertIsNone(w)
+
+    def test_union_type_positional_subpattern(self):
+        IntOrStr = int | str
+        w = None
+        match 0:
+            case IntOrStr(y):
+                w = y
+        self.assertEqual(w, 0)
+
+    def test_union_type_keyword_subpattern(self):
+        EitherPoint = Point | Point3D
+        p = Point(x=1, y=2)
+        w = None
+        match p:
+            case EitherPoint(x=1, y=2):
+                w = 0
+        self.assertEqual(w, 0)
+
+    def test_patma_union_arg(self):
+        p = Point(x=1, y=2)
+        IntOrStr = int | str
+        w = None
+        match p:
+            case Point(IntOrStr(), IntOrStr()):
+                w = 0
+        self.assertEqual(w, 0)
+
+    def test_patma_union_kwarg(self):
+        p = Point(x=1, y=2)
+        IntOrStr = int | str
+        w = None
+        match p:
+            case Point(x=IntOrStr(), y=IntOrStr()):
+                w = 0
+        self.assertEqual(w, 0)
+
+    def test_patma_union_arg_no_match(self):
+        p = Point(x=1, y=2)
+        StrOrBytes = str | bytes
+        w = None
+        match p:
+            case Point(StrOrBytes(), StrOrBytes()):
+                w = 0
+        self.assertIsNone(w)
+
+    def test_patma_union_kwarg_no_match(self):
+        p = Point(x=1, y=2)
+        StrOrBytes = str | bytes
+        w = None
+        match p:
+            case Point(x=StrOrBytes(), y=StrOrBytes()):
+                w = 0
+        self.assertIsNone(w)
+
+    def test_union_type_match_second_member(self):
+        EitherPoint = Point | Point3D
+        p = Point3D(x=1, y=2, z=3)
+        w = None
+        match p:
+            case EitherPoint(x=1, y=2, z=3):
+                w = 0
+        self.assertEqual(w, 0)
 
 
 class TestSyntaxErrors(unittest.TestCase):
@@ -3239,8 +3315,28 @@ def test_mapping_pattern_duplicate_key_edge_case3(self):
                 pass
         """)
 
+
 class TestTypeErrors(unittest.TestCase):
 
+    def test_generic_type(self):
+        t = list[str]
+        w = None
+        with self.assertRaises(TypeError):
+            match ["s"]:
+                case t():
+                    w = 0
+        self.assertIsNone(w)
+
+    def test_legacy_generic_type(self):
+        from typing import List
+        t = List[str]
+        w = None
+        with self.assertRaises(TypeError):
+            match ["s"]:
+                case t():
+                    w = 0
+        self.assertIsNone(w)
+
     def test_accepts_positional_subpatterns_0(self):
         class Class:
             __match_args__ = ()
@@ -3350,6 +3446,124 @@ def test_class_pattern_not_type(self):
                     w = 0
         self.assertIsNone(w)
 
+    def test_class_or_union_not_specialform(self):
+        from typing import Literal
+        name = type(Literal).__name__
+        msg = rf"called match pattern must be a class or types.UnionType \(got {name}\)"
+        w = None
+        with self.assertRaisesRegex(TypeError, msg):
+            match 1:
+                case Literal():
+                    w = 0
+        self.assertIsNone(w)
+
+    def test_legacy_union_type(self):
+        from typing import Union
+        IntOrStr = Union[int, str]
+        name = type(IntOrStr).__name__
+        msg = rf"called match pattern must be a class or types.UnionType \(got {name}\)"
+        w = None
+        with self.assertRaisesRegex(TypeError, msg):
+            match 1:
+                case IntOrStr():
+                    w = 0
+        self.assertIsNone(w)
+
+    def test_expanded_union_mirrors_isinstance_success(self):
+        ListOfInt = list[int]
+        t = int | ListOfInt
+        try:  # get the isinstance result
+            reference = isinstance(1, t)
+        except TypeError as exc:
+            reference = exc
+
+        try:  # get the match-case result
+            match 1:
+                case int() | ListOfInt():
+                    result = True
+                case _:
+                    result = False
+        except TypeError as exc:
+            result = exc
+
+        # we should ge the same result
+        self.assertIs(result, True)
+        self.assertIs(reference, True)
+
+    def test_expanded_union_mirrors_isinstance_failure(self):
+        ListOfInt = list[int]
+        t = ListOfInt | int
+
+        try:  # get the isinstance result
+            reference = isinstance(1, t)
+        except TypeError as exc:
+            reference = exc
+
+        try:  # get the match-case result
+            match 1:
+                case ListOfInt() | int():
+                    result = True
+                case _:
+                    result = False
+        except TypeError as exc:
+            result = exc
+
+        # we should ge the same result
+        self.assertIsInstance(result, TypeError)
+        self.assertIsInstance(reference, TypeError)
+
+    def test_union_type_mirrors_isinstance_success(self):
+        t = int | list[int]
+
+        try:  # get the isinstance result
+            reference = isinstance(1, t)
+        except TypeError as exc:
+            reference = exc
+
+        try:  # get the match-case result
+            match 1:
+                case t():
+                    result = True
+                case _:
+                    result = False
+        except TypeError as exc:
+            result = exc
+
+        # we should ge the same result
+        self.assertIs(result, True)
+        self.assertIs(reference, True)
+
+    def test_union_type_mirrors_isinstance_failure(self):
+        t = list[int] | int
+
+        try:  # get the isinstance result
+            reference = isinstance(1, t)
+        except TypeError as exc:
+            reference = exc
+
+        try:  # get the match-case result
+            match 1:
+                case t():
+                    result = True
+                case _:
+                    result = False
+        except TypeError as exc:
+            result = exc
+
+        # we should ge the same result
+        self.assertIsInstance(result, TypeError)
+        self.assertIsInstance(reference, TypeError)
+
+    def test_generic_union_type(self):
+        from collections.abc import Sequence, Set
+        t = Sequence[str] | Set[str]
+        w = None
+        with self.assertRaises(TypeError):
+            match ["s"]:
+                case t():
+                    w = 0
+        self.assertIsNone(w)
+
     def test_regular_protocol(self):
         from typing import Protocol
         class P(Protocol): ...
@@ -3379,31 +3593,6 @@ class A:
                     w = 0
         self.assertIsNone(w)
 
-    def test_union_type_postional_subpattern(self):
-        IntOrStr = int | str
-        x = 1
-        w = None
-        with self.assertRaises(TypeError):
-            match x:
-                case IntOrStr(x):
-                    w = 0
-        self.assertEqual(x, 1)
-        self.assertIsNone(w)
-
-    def test_union_type_keyword_subpattern(self):
-        @dataclasses.dataclass
-        class Point2:
-            x: int
-            y: int
-        EitherPoint = Point | Point2
-        x = Point(x=1, y=2)
-        w = None
-        with self.assertRaises(TypeError):
-            match x:
-                case EitherPoint(x=1, y=2):
-                    w = 0
-        self.assertIsNone(w)
-
 
 class TestValueErrors(unittest.TestCase):
 

Python/ceval.c

@@ -726,26 +726,34 @@ PyObject*
 _PyEval_MatchClass(PyThreadState *tstate, PyObject *subject, PyObject *type,
                    Py_ssize_t nargs, PyObject *kwargs)
 {
-    if (!PyType_Check(type) && !_PyUnion_Check(type)) {
-        const char *e = "called match pattern must be a class or a union";
-        _PyErr_Format(tstate, PyExc_TypeError, e);
+    // Recurse on unions.
+    if (_PyUnion_Check(type)) {
+        // get union members
+        PyObject *members = _Py_union_args(type);
+        const Py_ssize_t n = PyTuple_GET_SIZE(members);
+
+        // iterate over union members and return first match
+        for (Py_ssize_t i = 0; i < n; i++) {
+            PyObject *member = PyTuple_GET_ITEM(members, i);
+            PyObject *attrs = _PyEval_MatchClass(tstate, subject, member, nargs, kwargs);
+            // match found
+            if (attrs != NULL) {
+                return attrs;
+            }
+        }
+        // no match found
+        return NULL;
+    }
+    if (!PyType_Check(type)) {
+        const char *e = "called match pattern must be a class or types.UnionType (got %s)";
+        _PyErr_Format(tstate, PyExc_TypeError, e, Py_TYPE(type)->tp_name);
         return NULL;
     }
     assert(PyTuple_CheckExact(kwargs));
     // First, an isinstance check:
     if (PyObject_IsInstance(subject, type) <= 0) {
         return NULL;
     }
-    // Subpatterns are not supported for union types:
-    if (_PyUnion_Check(type)) {
-        // Return error if any positional or keyword arguments are given:
-        if (nargs || PyTuple_GET_SIZE(kwargs)) {
-            const char *e = "union types do not support sub-patterns";
-            _PyErr_Format(tstate, PyExc_TypeError, e);
-            return NULL;
-        }
-        return PyTuple_New(0);
-    }
     // So far so good:
     PyObject *seen = PySet_New(NULL);
     if (seen == NULL) {