bpo-45292: [PEP-654] add except* (GH-29581)

1 files changed, 135 insertions, 23 deletions

diff --git a/Objects/exceptions.c b/Objects/exceptions.c
index 1340157..1db49d9 100644
--- a/Objects/exceptions.c
+++ b/Objects/exceptions.c
@@ -773,6 +773,23 @@ error:
     return NULL;
 }
 
+PyObject *
+_PyExc_CreateExceptionGroup(const char *msg_str, PyObject *excs)
+{
+    PyObject *msg = PyUnicode_FromString(msg_str);
+    if (!msg) {
+        return NULL;
+    }
+    PyObject *args = PyTuple_Pack(2, msg, excs);
+    Py_DECREF(msg);
+    if (!args) {
+        return NULL;
+    }
+    PyObject *result = PyObject_CallObject(PyExc_BaseExceptionGroup, args);
+    Py_DECREF(args);
+    return result;
+}
+
 static int
 BaseExceptionGroup_init(PyBaseExceptionGroupObject *self,
     PyObject *args, PyObject *kwds)
@@ -878,7 +895,7 @@ exceptiongroup_subset(
     if (tb) {
         int res = PyException_SetTraceback(eg, tb);
         Py_DECREF(tb);
-        if (res == -1) {
+        if (res < 0) {
             goto error;
         }
     }
@@ -896,26 +913,41 @@ typedef enum {
     EXCEPTION_GROUP_MATCH_BY_TYPE = 0,
     /* A PyFunction returning True for matching exceptions */
     EXCEPTION_GROUP_MATCH_BY_PREDICATE = 1,
-    /* An instance or container thereof, checked with equality
-     * This matcher type is only used internally by the
-     * interpreter, it is not exposed to python code */
+    /* A set of leaf exceptions to include in the result.
+     * This matcher type is used internally by the interpreter
+     * to construct reraised exceptions.
+     */
     EXCEPTION_GROUP_MATCH_INSTANCES = 2
 } _exceptiongroup_split_matcher_type;
 
 static int
 get_matcher_type(PyObject *value,
-                  _exceptiongroup_split_matcher_type *type)
+                 _exceptiongroup_split_matcher_type *type)
 {
-    /* the python API supports only BY_TYPE and BY_PREDICATE */
-    if (PyExceptionClass_Check(value) ||
-        PyTuple_CheckExact(value)) {
+    assert(value);
+
+    if (PyFunction_Check(value)) {
+        *type = EXCEPTION_GROUP_MATCH_BY_PREDICATE;
+        return 0;
+    }
+
+    if (PyExceptionClass_Check(value)) {
         *type = EXCEPTION_GROUP_MATCH_BY_TYPE;
         return 0;
     }
-    if (PyFunction_Check(value)) {
-        *type = EXCEPTION_GROUP_MATCH_BY_PREDICATE;
+
+    if (PyTuple_CheckExact(value)) {
+        Py_ssize_t n = PyTuple_GET_SIZE(value);
+        for (Py_ssize_t i=0; i<n; i++) {
+            if (!PyExceptionClass_Check(PyTuple_GET_ITEM(value, i))) {
+                goto error;
+            }
+        }
+        *type = EXCEPTION_GROUP_MATCH_BY_TYPE;
         return 0;
     }
+
+error:
     PyErr_SetString(
         PyExc_TypeError,
         "expected a function, exception type or tuple of exception types");
@@ -944,10 +976,11 @@ exceptiongroup_split_check_match(PyObject *exc,
         return is_true;
     }
     case EXCEPTION_GROUP_MATCH_INSTANCES: {
-        if (PySequence_Check(matcher_value)) {
-            return PySequence_Contains(matcher_value, exc);
+        assert(PySet_Check(matcher_value));
+        if (!_PyBaseExceptionGroup_Check(exc)) {
+            return PySet_Contains(matcher_value, exc);
         }
-        return matcher_value == exc;
+        return 0;
     }
     }
     return 0;
@@ -1019,7 +1052,7 @@ exceptiongroup_split_recursive(PyObject *exc,
         }
         if (exceptiongroup_split_recursive(
                 e, matcher_type, matcher_value,
-                construct_rest, &rec_result) == -1) {
+                construct_rest, &rec_result) < 0) {
             assert(!rec_result.match);
             assert(!rec_result.rest);
             Py_LeaveRecursiveCall();
@@ -1028,7 +1061,7 @@ exceptiongroup_split_recursive(PyObject *exc,
         Py_LeaveRecursiveCall();
         if (rec_result.match) {
             assert(PyList_CheckExact(match_list));
-            if (PyList_Append(match_list, rec_result.match) == -1) {
+            if (PyList_Append(match_list, rec_result.match) < 0) {
                 Py_DECREF(rec_result.match);
                 goto done;
             }
@@ -1037,7 +1070,7 @@ exceptiongroup_split_recursive(PyObject *exc,
         if (rec_result.rest) {
             assert(construct_rest);
             assert(PyList_CheckExact(rest_list));
-            if (PyList_Append(rest_list, rec_result.rest) == -1) {
+            if (PyList_Append(rest_list, rec_result.rest) < 0) {
                 Py_DECREF(rec_result.rest);
                 goto done;
             }
@@ -1046,13 +1079,13 @@ exceptiongroup_split_recursive(PyObject *exc,
     }
 
     /* construct result */
-    if (exceptiongroup_subset(eg, match_list, &result->match) == -1) {
+    if (exceptiongroup_subset(eg, match_list, &result->match) < 0) {
         goto done;
     }
 
     if (construct_rest) {
         assert(PyList_CheckExact(rest_list));
-        if (exceptiongroup_subset(eg, rest_list, &result->rest) == -1) {
+        if (exceptiongroup_subset(eg, rest_list, &result->rest) < 0) {
             Py_CLEAR(result->match);
             goto done;
         }
@@ -1061,7 +1094,7 @@ exceptiongroup_split_recursive(PyObject *exc,
 done:
     Py_DECREF(match_list);
     Py_XDECREF(rest_list);
-    if (retval == -1) {
+    if (retval < 0) {
         Py_CLEAR(result->match);
         Py_CLEAR(result->rest);
     }
@@ -1077,7 +1110,7 @@ BaseExceptionGroup_split(PyObject *self, PyObject *args)
     }
 
     _exceptiongroup_split_matcher_type matcher_type;
-    if (get_matcher_type(matcher_value, &matcher_type) == -1) {
+    if (get_matcher_type(matcher_value, &matcher_type) < 0) {
         return NULL;
     }
 
@@ -1085,7 +1118,7 @@ BaseExceptionGroup_split(PyObject *self, PyObject *args)
     bool construct_rest = true;
     if (exceptiongroup_split_recursive(
             self, matcher_type, matcher_value,
-            construct_rest, &split_result) == -1) {
+            construct_rest, &split_result) < 0) {
         return NULL;
     }
 
@@ -1108,7 +1141,7 @@ BaseExceptionGroup_subgroup(PyObject *self, PyObject *args)
     }
 
     _exceptiongroup_split_matcher_type matcher_type;
-    if (get_matcher_type(matcher_value, &matcher_type) == -1) {
+    if (get_matcher_type(matcher_value, &matcher_type) < 0) {
         return NULL;
     }
 
@@ -1116,7 +1149,7 @@ BaseExceptionGroup_subgroup(PyObject *self, PyObject *args)
     bool construct_rest = false;
     if (exceptiongroup_split_recursive(
             self, matcher_type, matcher_value,
-            construct_rest, &split_result) == -1) {
+            construct_rest, &split_result) < 0) {
         return NULL;
     }
 
@@ -1128,6 +1161,85 @@ BaseExceptionGroup_subgroup(PyObject *self, PyObject *args)
     return result;
 }
 
+static int
+collect_exception_group_leaves(PyObject *exc, PyObject *leaves)
+{
+    if (Py_IsNone(exc)) {
+        return 0;
+    }
+
+    assert(PyExceptionInstance_Check(exc));
+    assert(PySet_Check(leaves));
+
+    /* Add all leaf exceptions in exc to the leaves set */
+
+    if (!_PyBaseExceptionGroup_Check(exc)) {
+        if (PySet_Add(leaves, exc) < 0) {
+            return -1;
+        }
+        return 0;
+    }
+    PyBaseExceptionGroupObject *eg = _PyBaseExceptionGroupObject_cast(exc);
+    Py_ssize_t num_excs = PyTuple_GET_SIZE(eg->excs);
+    /* recursive calls */
+    for (Py_ssize_t i = 0; i < num_excs; i++) {
+        PyObject *e = PyTuple_GET_ITEM(eg->excs, i);
+        if (Py_EnterRecursiveCall(" in collect_exception_group_leaves")) {
+            return -1;
+        }
+        int res = collect_exception_group_leaves(e, leaves);
+        Py_LeaveRecursiveCall();
+        if (res < 0) {
+            return -1;
+        }
+    }
+    return 0;
+}
+
+/* This function is used by the interpreter to construct reraised
+ * exception groups. It takes an exception group eg and a list
+ * of exception groups keep and returns the sub-exception group
+ * of eg which contains all leaf exceptions that are contained
+ * in any exception group in keep.
+ */
+PyObject *
+_PyExc_ExceptionGroupProjection(PyObject *eg, PyObject *keep)
+{
+    assert(_PyBaseExceptionGroup_Check(eg));
+    assert(PyList_CheckExact(keep));
+
+    PyObject *leaves = PySet_New(NULL);
+    if (!leaves) {
+        return NULL;
+    }
+
+    Py_ssize_t n = PyList_GET_SIZE(keep);
+    for (Py_ssize_t i = 0; i < n; i++) {
+        PyObject *e = PyList_GET_ITEM(keep, i);
+        assert(e != NULL);
+        assert(_PyBaseExceptionGroup_Check(e));
+        if (collect_exception_group_leaves(e, leaves) < 0) {
+            Py_DECREF(leaves);
+            return NULL;
+        }
+    }
+
+    _exceptiongroup_split_result split_result;
+    bool construct_rest = false;
+    int err = exceptiongroup_split_recursive(
+                eg, EXCEPTION_GROUP_MATCH_INSTANCES, leaves,
+                construct_rest, &split_result);
+    Py_DECREF(leaves);
+    if (err < 0) {
+        return NULL;
+    }
+
+    PyObject *result = split_result.match ?
+        split_result.match : Py_NewRef(Py_None);
+    assert(split_result.rest == NULL);
+    return result;
+}
+
 static PyMemberDef BaseExceptionGroup_members[] = {
     {"message", T_OBJECT, offsetof(PyBaseExceptionGroupObject, msg), READONLY,
         PyDoc_STR("exception message")},