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

1 files changed, 338 insertions, 1 deletions

diff --git a/Python/compile.c b/Python/compile.c
index bbdb011..00e1e01 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -175,7 +175,7 @@ compiler IR.
 
 enum fblocktype { WHILE_LOOP, FOR_LOOP, TRY_EXCEPT, FINALLY_TRY, FINALLY_END,
                   WITH, ASYNC_WITH, HANDLER_CLEANUP, POP_VALUE, EXCEPTION_HANDLER,
-                  ASYNC_COMPREHENSION_GENERATOR };
+                  EXCEPTION_GROUP_HANDLER, ASYNC_COMPREHENSION_GENERATOR };
 
 struct fblockinfo {
     enum fblocktype fb_type;
@@ -323,6 +323,7 @@ static int compiler_call_helper(struct compiler *c, int n,
                                 asdl_expr_seq *args,
                                 asdl_keyword_seq *keywords);
 static int compiler_try_except(struct compiler *, stmt_ty);
+static int compiler_try_star_except(struct compiler *, stmt_ty);
 static int compiler_set_qualname(struct compiler *);
 
 static int compiler_sync_comprehension_generator(
@@ -1094,6 +1095,8 @@ stack_effect(int opcode, int oparg, int jump)
             return -1;
         case JUMP_IF_NOT_EXC_MATCH:
             return -1;
+        case JUMP_IF_NOT_EG_MATCH:
+            return jump > 0 ? -1 : 2;
         case IMPORT_NAME:
             return -1;
         case IMPORT_FROM:
@@ -1131,6 +1134,8 @@ stack_effect(int opcode, int oparg, int jump)
              * if an exception be raised. */
             return jump ? -1 + 4 : 0;
 
+        case PREP_RERAISE_STAR:
+             return 2;
         case RERAISE:
             return -3;
         case PUSH_EXC_INFO:
@@ -1755,6 +1760,18 @@ find_ann(asdl_stmt_seq *stmts)
                   find_ann(st->v.Try.finalbody) ||
                   find_ann(st->v.Try.orelse);
             break;
+        case TryStar_kind:
+            for (j = 0; j < asdl_seq_LEN(st->v.TryStar.handlers); j++) {
+                excepthandler_ty handler = (excepthandler_ty)asdl_seq_GET(
+                    st->v.TryStar.handlers, j);
+                if (find_ann(handler->v.ExceptHandler.body)) {
+                    return 1;
+                }
+            }
+            res = find_ann(st->v.TryStar.body) ||
+                  find_ann(st->v.TryStar.finalbody) ||
+                  find_ann(st->v.TryStar.orelse);
+            break;
         default:
             res = 0;
         }
@@ -1816,6 +1833,7 @@ compiler_unwind_fblock(struct compiler *c, struct fblockinfo *info,
     switch (info->fb_type) {
         case WHILE_LOOP:
         case EXCEPTION_HANDLER:
+        case EXCEPTION_GROUP_HANDLER:
         case ASYNC_COMPREHENSION_GENERATOR:
             return 1;
 
@@ -1919,6 +1937,10 @@ compiler_unwind_fblock_stack(struct compiler *c, int preserve_tos, struct fblock
         return 1;
     }
     struct fblockinfo *top = &c->u->u_fblock[c->u->u_nfblocks-1];
+    if (top->fb_type == EXCEPTION_GROUP_HANDLER) {
+        return compiler_error(
+            c, "'break', 'continue' and 'return' cannot appear in an except* block");
+    }
     if (loop != NULL && (top->fb_type == WHILE_LOOP || top->fb_type == FOR_LOOP)) {
         *loop = top;
         return 1;
@@ -3202,6 +3224,62 @@ compiler_try_finally(struct compiler *c, stmt_ty s)
     return 1;
 }
 
+static int
+compiler_try_star_finally(struct compiler *c, stmt_ty s)
+{
+    basicblock *body = compiler_new_block(c);
+    if (body == NULL) {
+        return 0;
+    }
+    basicblock *end = compiler_new_block(c);
+    if (!end) {
+        return 0;
+    }
+    basicblock *exit = compiler_new_block(c);
+    if (!exit) {
+        return 0;
+    }
+    basicblock *cleanup = compiler_new_block(c);
+    if (!cleanup) {
+        return 0;
+    }
+    /* `try` block */
+    ADDOP_JUMP(c, SETUP_FINALLY, end);
+    compiler_use_next_block(c, body);
+    if (!compiler_push_fblock(c, FINALLY_TRY, body, end, s->v.TryStar.finalbody)) {
+        return 0;
+    }
+    if (s->v.TryStar.handlers && asdl_seq_LEN(s->v.TryStar.handlers)) {
+        if (!compiler_try_star_except(c, s)) {
+            return 0;
+        }
+    }
+    else {
+        VISIT_SEQ(c, stmt, s->v.TryStar.body);
+    }
+    ADDOP_NOLINE(c, POP_BLOCK);
+    compiler_pop_fblock(c, FINALLY_TRY, body);
+    VISIT_SEQ(c, stmt, s->v.TryStar.finalbody);
+    ADDOP_JUMP_NOLINE(c, JUMP_FORWARD, exit);
+    /* `finally` block */
+    compiler_use_next_block(c, end);
+
+    UNSET_LOC(c);
+    ADDOP_JUMP(c, SETUP_CLEANUP, cleanup);
+    ADDOP(c, PUSH_EXC_INFO);
+    if (!compiler_push_fblock(c, FINALLY_END, end, NULL, NULL)) {
+        return 0;
+    }
+    VISIT_SEQ(c, stmt, s->v.TryStar.finalbody);
+    compiler_pop_fblock(c, FINALLY_END, end);
+    ADDOP_I(c, RERAISE, 0);
+    compiler_use_next_block(c, cleanup);
+    ADDOP(c, POP_EXCEPT_AND_RERAISE);
+    compiler_use_next_block(c, exit);
+    return 1;
+}
+
+
 /*
    Code generated for "try: S except E1 as V1: S1 except E2 as V2: S2 ...":
    (The contents of the value stack is shown in [], with the top
@@ -3360,6 +3438,253 @@ compiler_try_except(struct compiler *c, stmt_ty s)
     ADDOP(c, POP_EXCEPT_AND_RERAISE);
     compiler_use_next_block(c, orelse);
     VISIT_SEQ(c, stmt, s->v.Try.orelse);
+    ADDOP_JUMP(c, JUMP_FORWARD, end);
+    compiler_use_next_block(c, end);
+    return 1;
+}
+
+/*
+   Code generated for "try: S except* E1 as V1: S1 except* E2 as V2: S2 ...":
+   (The contents of the value stack is shown in [], with the top
+   at the right; 'tb' is trace-back info, 'val' the exception instance,
+   and 'typ' the exception's type.)
+
+   Value stack                             Label         Instruction     Argument
+   []                                                   SETUP_FINALLY         L1
+   []                                                   <code for S>
+   []                                                   POP_BLOCK
+   []                                                   JUMP_FORWARD          L0
+
+   [tb, val, typ]                            L1:        DUP_TOP_TWO               )  save a copy of the
+   [tb, val, typ, orig, typ]                            POP_TOP                   )  original raised exception
+   [tb, val, typ, orig]                                 ROT_FOUR                  )
+   [orig, tb, val, typ]                                 BUILD_LIST                )  list for raised/reraised
+   [orig, tb, val, typ, res]                            ROT_FOUR                  )  exceptions ("result")
+
+   [orig, res, tb, val, typ]                            <evaluate E1>             )
+   [orig, res, tb, val, typ, E1]                        JUMP_IF_NOT_EXC_MATCH L2  ) only if E1
+
+   [orig, res, tb, rest, typ, tb, match, typ]           POP
+   [orig, res, tb, rest, typ, tb, match]                <assign to V1>  (or POP if no V1)
+   [orig, res, tb, rest, typ, tb]                       POP
+
+   [orig, res, tb, rest, typ]                           SETUP_FINALLY         R1
+   [orig, res, tb, rest, typ]                           <code for S1>
+   [orig, res, tb, rest, typ]                           JUMP_FORWARD          L2
+
+   [orig, res, tb, rest, typ, i, tb, v, t]   R1:        POP           ) exception raised in except* body
+   [orig, res, tb, rest, typ, i, tb, v]                 LIST_APPEND 6 ) add it to res
+   [orig, res, tb, rest, typ, i, tb]                    POP
+   [orig, res, tb, rest, typ, i]                        POP
+
+   [orig, res, tb, rest, typ]                L2:        <evaluate E2>
+   .............................etc.......................
+
+   [orig, res, tb, rest, typ]                Ln+1:       POP           ) add unhandled exception
+   [orig, res, tb, rest]                                 LIST_APPEND 2 ) to res (could be None)
+   [orig, res, tb]                                       POP
+
+   [orig, res]                                           PREP_RERAISE_STAR
+   [i, tb, val, typ]                                     POP_JUMP_IF_TRUE      RER
+   [i, tb, val, typ]                                     POP
+   [i, tb, val]                                          POP
+   [i, tb]                                               POP
+   [i]                                                   POP
+   []                                                    JUMP_FORWARD          L0
+
+   [i, tb, val, typ]                        RER:         POP_EXCEPT_AND_RERAISE
+
+   []                                       L0:     <next statement>
+*/
+static int
+compiler_try_star_except(struct compiler *c, stmt_ty s)
+{
+    basicblock *body = compiler_new_block(c);
+    if (body == NULL) {
+        return 0;
+    }
+    basicblock *except = compiler_new_block(c);
+    if (except == NULL) {
+        return 0;
+    }
+    basicblock *orelse = compiler_new_block(c);
+     if (orelse == NULL) {
+        return 0;
+    }
+    basicblock *end = compiler_new_block(c);
+    if (end == NULL) {
+        return 0;
+    }
+    basicblock *cleanup = compiler_new_block(c);
+    if (cleanup == NULL) {
+        return 0;
+    }
+    basicblock *reraise_star = compiler_new_block(c);
+    if (reraise_star == NULL) {
+        return 0;
+    }
+
+    ADDOP_JUMP(c, SETUP_FINALLY, except);
+    compiler_use_next_block(c, body);
+    if (!compiler_push_fblock(c, TRY_EXCEPT, body, NULL, NULL)) {
+        return 0;
+    }
+    VISIT_SEQ(c, stmt, s->v.TryStar.body);
+    compiler_pop_fblock(c, TRY_EXCEPT, body);
+    ADDOP_NOLINE(c, POP_BLOCK);
+    ADDOP_JUMP_NOLINE(c, JUMP_FORWARD, orelse);
+    Py_ssize_t n = asdl_seq_LEN(s->v.TryStar.handlers);
+    compiler_use_next_block(c, except);
+
+    UNSET_LOC(c);
+    ADDOP_JUMP(c, SETUP_CLEANUP, cleanup);
+    ADDOP(c, PUSH_EXC_INFO);
+    /* Runtime will push a block here, so we need to account for that */
+    if (!compiler_push_fblock(c, EXCEPTION_GROUP_HANDLER,
+                                 NULL, NULL, "except handler")) {
+        return 0;
+    }
+    for (Py_ssize_t i = 0; i < n; i++) {
+        excepthandler_ty handler = (excepthandler_ty)asdl_seq_GET(
+            s->v.TryStar.handlers, i);
+        SET_LOC(c, handler);
+        except = compiler_new_block(c);
+        if (except == NULL) {
+            return 0;
+        }
+        if (i == 0) {
+            /* Push the original EG into the stack */
+            /*
+               [tb, val, exc]            DUP_TOP_TWO
+               [tb, val, exc, val, exc]  POP_TOP
+               [tb, val, exc, val]       ROT_FOUR
+               [val, tb, val, exc]
+            */
+            ADDOP(c, DUP_TOP_TWO);
+            ADDOP(c, POP_TOP);
+            ADDOP(c, ROT_FOUR);
+
+            /* create empty list for exceptions raised/reraise in the except* blocks */
+            /*
+               [val, tb, val, exc]       BUILD_LIST
+               [val, tb, val, exc, []]   ROT_FOUR
+               [val, [], tb, val, exc]
+            */
+            ADDOP_I(c, BUILD_LIST, 0);
+            ADDOP(c, ROT_FOUR);
+        }
+        if (handler->v.ExceptHandler.type) {
+            VISIT(c, expr, handler->v.ExceptHandler.type);
+            ADDOP_JUMP(c, JUMP_IF_NOT_EG_MATCH, except);
+            NEXT_BLOCK(c);
+        }
+        ADDOP(c, POP_TOP);  // exc_type
+
+        basicblock *cleanup_end = compiler_new_block(c);
+        if (cleanup_end == NULL) {
+            return 0;
+        }
+        basicblock *cleanup_body = compiler_new_block(c);
+        if (cleanup_body == NULL) {
+            return 0;
+        }
+
+        if (handler->v.ExceptHandler.name) {
+            compiler_nameop(c, handler->v.ExceptHandler.name, Store);
+        }
+        else {
+            ADDOP(c, POP_TOP);  // val
+        }
+        ADDOP(c, POP_TOP);  // tb
+
+        /*
+          try:
+              # body
+          except type as name:
+              try:
+                  # body
+              finally:
+                  name = None # in case body contains "del name"
+                  del name
+        */
+        /* second try: */
+        ADDOP_JUMP(c, SETUP_CLEANUP, cleanup_end);
+        compiler_use_next_block(c, cleanup_body);
+        if (!compiler_push_fblock(c, HANDLER_CLEANUP, cleanup_body, NULL, handler->v.ExceptHandler.name))
+            return 0;
+
+        /* second # body */
+        VISIT_SEQ(c, stmt, handler->v.ExceptHandler.body);
+        compiler_pop_fblock(c, HANDLER_CLEANUP, cleanup_body);
+        /* name = None; del name; # Mark as artificial */
+        UNSET_LOC(c);
+        ADDOP(c, POP_BLOCK);
+        if (handler->v.ExceptHandler.name) {
+            ADDOP_LOAD_CONST(c, Py_None);
+            compiler_nameop(c, handler->v.ExceptHandler.name, Store);
+            compiler_nameop(c, handler->v.ExceptHandler.name, Del);
+        }
+        ADDOP_JUMP(c, JUMP_FORWARD, except);
+
+        /* except: */
+        compiler_use_next_block(c, cleanup_end);
+
+        /* name = None; del name; # Mark as artificial */
+        UNSET_LOC(c);
+
+        if (handler->v.ExceptHandler.name) {
+            ADDOP_LOAD_CONST(c, Py_None);
+            compiler_nameop(c, handler->v.ExceptHandler.name, Store);
+            compiler_nameop(c, handler->v.ExceptHandler.name, Del);
+        }
+
+        /* add exception raised to the res list */
+        ADDOP(c, POP_TOP); // type
+        ADDOP_I(c, LIST_APPEND, 6); // exc
+        ADDOP(c, POP_TOP); // tb
+        ADDOP(c, POP_TOP); // lasti
+
+        ADDOP_JUMP(c, JUMP_ABSOLUTE, except);
+        compiler_use_next_block(c, except);
+
+        if (i == n - 1) {
+            /* Add exc to the list (if not None it's the unhandled part of the EG) */
+            ADDOP(c, POP_TOP);
+            ADDOP_I(c, LIST_APPEND, 2);
+            ADDOP(c, POP_TOP);
+            ADDOP_JUMP(c, JUMP_FORWARD, reraise_star);
+        }
+    }
+    /* Mark as artificial */
+    UNSET_LOC(c);
+    compiler_pop_fblock(c, EXCEPTION_GROUP_HANDLER, NULL);
+    basicblock *reraise = compiler_new_block(c);
+    if (!reraise) {
+        return 0;
+    }
+
+    compiler_use_next_block(c, reraise_star);
+    ADDOP(c, PREP_RERAISE_STAR);
+    ADDOP(c, DUP_TOP);
+    ADDOP_JUMP(c, POP_JUMP_IF_TRUE, reraise);
+    NEXT_BLOCK(c);
+
+    /* Nothing to reraise - pop it */
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_TOP);
+    ADDOP(c, POP_BLOCK);
+    ADDOP(c, POP_EXCEPT);
+    ADDOP_JUMP(c, JUMP_FORWARD, end);
+    compiler_use_next_block(c, reraise);
+    ADDOP(c, POP_BLOCK);
+    ADDOP(c, POP_EXCEPT_AND_RERAISE);
+    compiler_use_next_block(c, cleanup);
+    ADDOP(c, POP_EXCEPT_AND_RERAISE);
+    compiler_use_next_block(c, orelse);
+    VISIT_SEQ(c, stmt, s->v.TryStar.orelse);
+    ADDOP_JUMP(c, JUMP_FORWARD, end);
     compiler_use_next_block(c, end);
     return 1;
 }
@@ -3372,6 +3697,16 @@ compiler_try(struct compiler *c, stmt_ty s) {
         return compiler_try_except(c, s);
 }
 
+static int
+compiler_try_star(struct compiler *c, stmt_ty s)
+{
+    if (s->v.TryStar.finalbody && asdl_seq_LEN(s->v.TryStar.finalbody)) {
+        return compiler_try_star_finally(c, s);
+    }
+    else {
+        return compiler_try_star_except(c, s);
+    }
+}
 
 static int
 compiler_import_as(struct compiler *c, identifier name, identifier asname)
@@ -3634,6 +3969,8 @@ compiler_visit_stmt(struct compiler *c, stmt_ty s)
         break;
     case Try_kind:
         return compiler_try(c, s);
+    case TryStar_kind:
+        return compiler_try_star(c, s);
     case Assert_kind:
         return compiler_assert(c, s);
     case Import_kind: