Add support for split() and subgroup() to BaseException.

Author: picnixz

Lib/test/test_exception_group.py

@@ -308,32 +308,39 @@ class ExceptionGroupSubgroupTests(ExceptionGroupTestBase):
     def setUp(self):
         self.eg = create_simple_eg()
         self.eg_template = [ValueError(1), TypeError(int), ValueError(2)]
+        self.exc = ValueError('a')
 
     def test_basics_subgroup_split__bad_arg_type(self):
-        class C:
-            pass
-
-        bad_args = ["bad arg",
-                    C,
-                    OSError('instance not type'),
-                    [OSError, TypeError],
-                    (OSError, 42),
-                   ]
-        for arg in bad_args:
-            with self.assertRaises(TypeError):
-                self.eg.subgroup(arg)
-            with self.assertRaises(TypeError):
-                self.eg.split(arg)
+        for obj in (self.eg, self.exc):
+            for arg in (
+                "bad arg",
+                type('NewClass', (), {}),
+                OSError('instance not type'),
+                [OSError, TypeError],
+                (OSError, 42),
+            ):
+                with self.subTest(obj_type=type(obj), arg=arg):
+                    with self.assertRaises(TypeError):
+                        obj.subgroup(arg)
+
+                    with self.assertRaises(TypeError):
+                        obj.split(arg)
 
     def test_basics_subgroup_by_type__passthrough(self):
-        eg = self.eg
-        self.assertIs(eg, eg.subgroup(BaseException))
-        self.assertIs(eg, eg.subgroup(Exception))
-        self.assertIs(eg, eg.subgroup(BaseExceptionGroup))
-        self.assertIs(eg, eg.subgroup(ExceptionGroup))
+        for exc_type in (
+            BaseException, Exception,
+            BaseExceptionGroup, ExceptionGroup,
+        ):
+            with self.subTest(exc_type):
+                self.assertIs(self.eg, self.eg.subgroup(exc_type))
+
+                wrapped = self.exc.subgroup(ValueError)
+                self.assertEqual(wrapped.message, str(self.exc))
+                self.assertMatchesTemplate(wrapped, ExceptionGroup, [self.exc])
 
     def test_basics_subgroup_by_type__no_match(self):
         self.assertIsNone(self.eg.subgroup(OSError))
+        self.assertIsNone(self.exc.subgroup(OSError))
 
     def test_basics_subgroup_by_type__match(self):
         eg = self.eg
@@ -349,15 +356,24 @@ def test_basics_subgroup_by_type__match(self):
                 self.assertEqual(subeg.message, eg.message)
                 self.assertMatchesTemplate(subeg, ExceptionGroup, template)
 
+        wrapped = self.exc.subgroup(ValueError)
+        self.assertEqual(wrapped.message, str(self.exc))
+        self.assertMatchesTemplate(wrapped, ExceptionGroup, [self.exc])
+
     def test_basics_subgroup_by_predicate__passthrough(self):
         f = lambda e: True
         for callable in [f, Predicate(f), Predicate(f).method]:
             self.assertIs(self.eg, self.eg.subgroup(callable))
 
+            wrapped = self.exc.subgroup(callable)
+            self.assertEqual(wrapped.message, str(self.exc))
+            self.assertMatchesTemplate(wrapped, ExceptionGroup, [self.exc])
+
     def test_basics_subgroup_by_predicate__no_match(self):
         f = lambda e: False
         for callable in [f, Predicate(f), Predicate(f).method]:
             self.assertIsNone(self.eg.subgroup(callable))
+            self.assertIsNone(self.exc.subgroup(callable))
 
     def test_basics_subgroup_by_predicate__match(self):
         eg = self.eg
@@ -371,40 +387,61 @@ def test_basics_subgroup_by_predicate__match(self):
             f = lambda e: isinstance(e, match_type)
             for callable in [f, Predicate(f), Predicate(f).method]:
                 with self.subTest(callable=callable):
-                    subeg = eg.subgroup(f)
+                    subeg = eg.subgroup(callable)
                     self.assertEqual(subeg.message, eg.message)
                     self.assertMatchesTemplate(subeg, ExceptionGroup, template)
 
+        f = lambda e: isinstance(e, ValueError)
+        for callable in [f, Predicate(f), Predicate(f).method]:
+            group = self.exc.subgroup(callable)
+            self.assertEqual(group.message, str(self.exc))
+            self.assertMatchesTemplate(group, ExceptionGroup, [self.exc])
+
 
 class ExceptionGroupSplitTests(ExceptionGroupTestBase):
     def setUp(self):
         self.eg = create_simple_eg()
         self.eg_template = [ValueError(1), TypeError(int), ValueError(2)]
 
+        self.exc = ValueError('a')
+
     def test_basics_split_by_type__passthrough(self):
-        for E in [BaseException, Exception,
-                  BaseExceptionGroup, ExceptionGroup]:
-            match, rest = self.eg.split(E)
-            self.assertMatchesTemplate(
-                match, ExceptionGroup, self.eg_template)
-            self.assertIsNone(rest)
+        for exc_type in (
+            BaseException, Exception,
+            BaseExceptionGroup, ExceptionGroup,
+        ):
+            with self.subTest(exc_type):
+                match, rest = self.eg.split(exc_type)
+                self.assertMatchesTemplate(match, ExceptionGroup,
+                                           self.eg_template)
+                self.assertIsNone(rest)
+
+                match, rest = self.exc.split(exc_type)
+                self.assertEqual(match.message, str(self.exc))
+                self.assertMatchesTemplate(match, ExceptionGroup, [self.exc])
+                self.assertIsNone(rest)
 
     def test_basics_split_by_type__no_match(self):
         match, rest = self.eg.split(OSError)
         self.assertIsNone(match)
-        self.assertMatchesTemplate(
-            rest, ExceptionGroup, self.eg_template)
+        self.assertMatchesTemplate(rest, ExceptionGroup, self.eg_template)
+
+        match, rest = self.exc.split(OSError)
+        self.assertIsNone(match)
+        self.assertMatchesTemplate(rest, ExceptionGroup, [self.exc])
 
     def test_basics_split_by_type__match(self):
         eg = self.eg
-        VE = ValueError
-        TE = TypeError
         testcases = [
-            # (matcher, match_template, rest_template)
-            (VE, [VE(1), VE(2)], [TE(int)]),
-            (TE, [TE(int)], [VE(1), VE(2)]),
-            ((VE, TE), self.eg_template, None),
-            ((OSError, VE), [VE(1), VE(2)], [TE(int)]),
+            # (exc_or_eg, matcher, match_template, rest_template)
+            (ValueError, [ValueError(1), ValueError(2)], [TypeError(int)]),
+            (TypeError, [TypeError(int)], [ValueError(1), ValueError(2)]),
+            ((ValueError, TypeError), self.eg_template, None),
+            (
+                (OSError, ValueError),
+                [ValueError(1), ValueError(2)],
+                [TypeError(int)],
+            ),
         ]
 
         for match_type, match_template, rest_template in testcases:
@@ -419,29 +456,41 @@ def test_basics_split_by_type__match(self):
             else:
                 self.assertIsNone(rest)
 
+        match, rest = self.exc.split(ValueError)
+        self.assertEqual(match.message, str(self.exc))
+        self.assertMatchesTemplate(match, ExceptionGroup, [self.exc])
+        self.assertIsNone(rest)
+
     def test_basics_split_by_predicate__passthrough(self):
         f = lambda e: True
         for callable in [f, Predicate(f), Predicate(f).method]:
             match, rest = self.eg.split(callable)
             self.assertMatchesTemplate(match, ExceptionGroup, self.eg_template)
             self.assertIsNone(rest)
 
+            match, rest = self.exc.split(callable)
+            self.assertEqual(match.message, str(self.exc))
+            self.assertMatchesTemplate(match, ExceptionGroup, [self.exc])
+            self.assertIsNone(rest)
+
     def test_basics_split_by_predicate__no_match(self):
         f = lambda e: False
         for callable in [f, Predicate(f), Predicate(f).method]:
             match, rest = self.eg.split(callable)
             self.assertIsNone(match)
             self.assertMatchesTemplate(rest, ExceptionGroup, self.eg_template)
 
+            match, rest = self.exc.split(callable)
+            self.assertIsNone(match)
+            self.assertMatchesTemplate(rest, ExceptionGroup, [self.exc])
+
     def test_basics_split_by_predicate__match(self):
         eg = self.eg
-        VE = ValueError
-        TE = TypeError
         testcases = [
             # (matcher, match_template, rest_template)
-            (VE, [VE(1), VE(2)], [TE(int)]),
-            (TE, [TE(int)], [VE(1), VE(2)]),
-            ((VE, TE), self.eg_template, None),
+            (ValueError, [ValueError(1), ValueError(2)], [TypeError(int)]),
+            (TypeError, [TypeError(int)], [ValueError(1), ValueError(2)]),
+            ((ValueError, TypeError), self.eg_template, None),
         ]
 
         for match_type, match_template, rest_template in testcases:
@@ -456,6 +505,13 @@ def test_basics_split_by_predicate__match(self):
                     self.assertMatchesTemplate(
                         rest, ExceptionGroup, rest_template)
 
+        f = lambda e: isinstance(e, ValueError)
+        for callable in [f, Predicate(f), Predicate(f).method]:
+            match, rest = self.exc.split(callable)
+            self.assertEqual(match.message, str(self.exc))
+            self.assertMatchesTemplate(match, ExceptionGroup, [self.exc])
+            self.assertIsNone(rest)
+
 
 class DeepRecursionInSplitAndSubgroup(unittest.TestCase):
     def make_deep_eg(self):

Objects/exceptions.c

@@ -269,18 +269,77 @@ BaseException_add_note(PyObject *self, PyObject *note)
     Py_RETURN_NONE;
 }
 
+/*
+ * Return an exception group wrapping 'self'.
+ *
+ * If 'self' is an exception group, a strong
+ * reference to 'self' is returned instead.
+ */
+static PyObject *
+base_exception_as_group(PyObject *self)
+{
+    if (!PyExceptionInstance_Check(self)) {
+        PyErr_BadInternalCall();
+        return NULL;
+    }
+
+    if (_PyBaseExceptionGroup_Check(self)) {
+        return Py_NewRef(self);
+    }
+
+    PyObject *message = PyObject_Str(self);
+    if (message == NULL) {
+        return NULL;
+    }
+    PyObject *wrapped = PyTuple_New(1);
+    if (wrapped == NULL) {
+        Py_DECREF(message);
+        return NULL;
+    }
+    PyTuple_SET_ITEM(wrapped, 0, Py_NewRef(self));
+    PyObject *group = PyObject_CallFunctionObjArgs(PyExc_BaseExceptionGroup,
+                                                   message, wrapped, NULL);
+    Py_DECREF(wrapped);
+    Py_DECREF(message);
+    return group;
+}
+
+static PyObject *
+BaseException_subgroup(PyObject *self, PyObject *matcher)
+{
+    PyObject *group = base_exception_as_group(self);
+    if (group == NULL) {
+        return NULL;
+    }
+    PyObject *res = PyObject_CallMethodOneArg(group, &_Py_ID(subgroup), matcher);
+    Py_DECREF(group);
+    return res;
+}
+
+static PyObject *
+BaseException_split(PyObject *self, PyObject *matcher)
+{
+    PyObject *group = base_exception_as_group(self);
+    if (group == NULL) {
+        return NULL;
+    }
+    PyObject *res = PyObject_CallMethodOneArg(group, &_Py_ID(split), matcher);
+    Py_DECREF(group);
+    return res;
+}
+
 PyDoc_STRVAR(add_note_doc,
 "Exception.add_note(note) --\n\
     add a note to the exception");
 
 static PyMethodDef BaseException_methods[] = {
-   {"__reduce__", (PyCFunction)BaseException_reduce, METH_NOARGS },
-   {"__setstate__", (PyCFunction)BaseException_setstate, METH_O },
-   {"with_traceback", (PyCFunction)BaseException_with_traceback, METH_O,
-    with_traceback_doc},
-   {"add_note", (PyCFunction)BaseException_add_note, METH_O,
-    add_note_doc},
-   {NULL, NULL, 0, NULL},
+    {"__reduce__", (PyCFunction)BaseException_reduce, METH_NOARGS, NULL},
+    {"__setstate__", BaseException_setstate, METH_O, NULL},
+    {"with_traceback", BaseException_with_traceback, METH_O, with_traceback_doc},
+    {"add_note", BaseException_add_note, METH_O, add_note_doc},
+    {"split", BaseException_split, METH_O, NULL},
+    {"subgroup", BaseException_subgroup, METH_O, NULL},
+    {NULL, NULL, 0, NULL},
 };
 
 static PyObject *