Allow the use of unions as match patterns

Author: tmke8

Lib/test/test_patma.py

@@ -4,6 +4,7 @@
 import dis
 import enum
 import inspect
+from re import I
 import sys
 import unittest
 from test import support
@@ -2888,6 +2889,14 @@ class B(A): ...
                         h = 1
                 self.assertEqual(h, 1)
 
+    def test_patma_union_type(self):
+        IntOrStr = int | str
+        x = 0
+        match x:
+            case IntOrStr():
+                x = 1
+        self.assertEqual(x, 1)
+
 
 class TestSyntaxErrors(unittest.TestCase):
 
@@ -3370,6 +3379,31 @@ 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

@@ -39,6 +39,7 @@
 #include "pycore_template.h"      // _PyTemplate_Build()
 #include "pycore_traceback.h"     // _PyTraceBack_FromFrame
 #include "pycore_tuple.h"         // _PyTuple_ITEMS()
+#include "pycore_unionobject.h"   // _PyUnion_Check()
 #include "pycore_uop_ids.h"       // Uops
 
 #include "dictobject.h"
@@ -725,8 +726,8 @@ PyObject*
 _PyEval_MatchClass(PyThreadState *tstate, PyObject *subject, PyObject *type,
                    Py_ssize_t nargs, PyObject *kwargs)
 {
-    if (!PyType_Check(type)) {
-        const char *e = "called match pattern must be a class";
+    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);
         return NULL;
     }
@@ -735,6 +736,16 @@ _PyEval_MatchClass(PyThreadState *tstate, PyObject *subject, PyObject *type,
     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) {