From 33b03a5ba6df6f100b3b2192db374ac1fb986c5b Mon Sep 17 00:00:00 2001
From: CF Bolz-Tereick <cfbolz@gmx.de>
Date: Wed, 25 Sep 2024 13:17:53 +0200
Subject: [PATCH] implement finally with exception groups

---
 pypy/interpreter/astcompiler/codegen.py       | 30 ++++++++++-----
 .../test/apptest_exceptiongroup.py            | 37 +++++++++++++++++++
 pypy/interpreter/pyopcode.py                  |  4 +-
 3 files changed, 61 insertions(+), 10 deletions(-)

diff --git a/pypy/interpreter/astcompiler/codegen.py b/pypy/interpreter/astcompiler/codegen.py
index 2e4682469be..551c2e053db 100644
--- a/pypy/interpreter/astcompiler/codegen.py
+++ b/pypy/interpreter/astcompiler/codegen.py
@@ -963,7 +963,7 @@ def _visit_try_except(self, tr):
         self._visit_body(tr.orelse)
         self.use_next_block(end)
 
-    def _visit_try_finally(self, tr):
+    def _visit_try_finally(self, tr, has_handlers, body, finalbody):
         body = self.new_block()
         end = self.new_block()
         exit = self.new_block()
@@ -972,22 +972,26 @@ def _visit_try_finally(self, tr):
         self.emit_jump(ops.SETUP_FINALLY, end)
         self.use_next_block(body)
         self.push_frame_block(F_FINALLY_TRY, body, end, tr)
-        if tr.handlers:
-            self._visit_try_except(tr)
+        if has_handlers:
+            if isinstance(tr, ast.Try):
+                self._visit_try_except(tr)
+            else:
+                assert isinstance(tr, ast.TryStar)
+                self._visit_try_except_star(tr)
         else:
-            self._visit_body(tr.body)
+            self._visit_body(body)
         self.no_position_info()
         self.emit_op(ops.POP_BLOCK)
 
         # finally block, unexceptional case
         self.pop_frame_block(F_FINALLY_TRY, body)
-        self._visit_body(tr.finalbody)
+        self._visit_body(finalbody)
         self.emit_jump(ops.JUMP_FORWARD, exit)
 
         # finally block, exceptional case
         self.use_next_block(end)
         self.push_frame_block(F_FINALLY_END, end)
-        self._visit_body(tr.finalbody)
+        self._visit_body(finalbody)
         self.pop_frame_block(F_FINALLY_END, end)
 
         # the RERAISE will be duplicated by duplicate_exits_without_lineno
@@ -998,11 +1002,19 @@ def _visit_try_finally(self, tr):
 
     def visit_Try(self, tr):
         if tr.finalbody:
-            return self._visit_try_finally(tr)
+            return self._visit_try_finally(
+                    tr, tr.handlers is not None, tr.body, tr.finalbody)
         else:
             return self._visit_try_except(tr)
 
     def visit_TryStar(self, tr):
+        if tr.finalbody:
+            return self._visit_try_finally(
+                    tr, tr.handlers is not None, tr.body, tr.finalbody)
+        else:
+            return self._visit_try_except_star(tr)
+
+    def _visit_try_except_star(self, tr):
         """
         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
@@ -3000,9 +3012,9 @@ def view(startblock):
             str_instr = "%5s: %s" % (instr.position_info[0], ops.opname[instr.opcode])
             if instr.opcode >= ops.HAVE_ARGUMENT and instr.jump is None:
                 str_instr += " %s" % (instr.arg, )
-            if hasattr(instr, '_stack_depth_after'):
+            if instr._stack_depth_after != -99:
                 str_instr += " stack depth after: %s" % instr._stack_depth_after
-                if instr._stack_depth_after < 0 and instr._stack_depth_after != -99:
+                if instr._stack_depth_after < 0:
                     fillcolor = "red"
             label.append(str_instr)
             if instr.jump is not None:
diff --git a/pypy/interpreter/astcompiler/test/apptest_exceptiongroup.py b/pypy/interpreter/astcompiler/test/apptest_exceptiongroup.py
index cc5346c13d1..a29b42306aa 100644
--- a/pypy/interpreter/astcompiler/test/apptest_exceptiongroup.py
+++ b/pypy/interpreter/astcompiler/test/apptest_exceptiongroup.py
@@ -113,3 +113,40 @@ def try_except_star_with_else_direct_return(x):
         with raises(SyntaxError) as info:
             exec(src)
         assert str(info.value).startswith(f"'{kw}' cannot appear in an except* block")
+
+def maybe_raise(err):
+    if err:
+        raise err
+    
+
+def with_finally(l, err):
+    try:
+        maybe_raise(err)
+    except* TypeError:
+        l.append(1)
+    except* ValueError:
+        l.append(2)
+    else:
+        l.append(3)
+    finally:
+        l.append(4)
+
+def test_finally():
+    l = []
+    with_finally(l, None)
+    assert l == [3, 4]
+    l = []
+    with_finally(l, ValueError())
+    assert l == [2, 4]
+    l = []
+    with_finally(l, TypeError())
+    assert l == [1, 4]
+    l = []
+    with_finally(l, ExceptionGroup('abc', [ValueError(), TypeError()]))
+    assert l == [1, 2, 4]
+    with raises(ZeroDivisionError):
+        l = []
+        with_finally(l, ZeroDivisionError())
+    assert l == [4]
+
+
diff --git a/pypy/interpreter/pyopcode.py b/pypy/interpreter/pyopcode.py
index c359d647f39..04f66250eb6 100644
--- a/pypy/interpreter/pyopcode.py
+++ b/pypy/interpreter/pyopcode.py
@@ -2167,10 +2167,12 @@ def exception_group_match(space, w_eg, w_typ):
         w_list = space.newlist([w_eg])
         w_wrapped = space.call_function(w_ExceptionGroup, space.newtext(''), w_list)
         return w_wrapped, space.w_None
-    else:
+    elif space.isinstance_w(w_eg, w_BaseExceptionGroup):
         w_tup = space.call_method(w_eg, 'split', w_typ)
         w_match, w_rest = space.unpackiterable(w_tup, 2)
         return w_match, w_rest
+    else:
+        return space.w_None, space.w_None
 
 ### helpers written at the application-level ###
 # Some of these functions are expected to be generally useful if other