added subpattern support

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):
@@ -2897,6 +2903,60 @@ def test_patma_union_type(self):
                 x = 1
         self.assertEqual(x, 1)
 
+    def test_union_type_positional_subpattern(self):
+        IntOrStr = int | str
+        x = 1
+        w = None
+        match x:
+            case IntOrStr(y):
+                w = y
+        self.assertEqual(w, 1)
+
+    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 = 1
+        self.assertEqual(w, 1)
+
+    def test_patma_union_no_match(self):
+        IntOrStr = int | str
+        x = None
+        match x:
+            case IntOrStr():
+                x = 1
+        self.assertIsNone(x)
+
+    def test_patma_union_arg(self):
+        p = Point(x=1, y=2)
+        IntOrStr = int | str
+        w = None
+        match p:
+            case Point(IntOrStr(), IntOrStr()):
+                w = 1
+        self.assertEqual(w, 1)
+
+    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 = 1
+        self.assertEqual(w, 1)
+
+    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 = 1
+        self.assertEqual(w, 1)
+
+
 
 class TestSyntaxErrors(unittest.TestCase):
 
@@ -3379,31 +3439,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);
+        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 a union, (got %s, %s)";
+        _PyErr_Format(tstate, PyExc_TypeError, e, Py_TYPE(type)->tp_name, subject->ob_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) {