From 3ee41410aefb44e32516c39b85513b37c66d214f Mon Sep 17 00:00:00 2001
From: Xinlong Hu <118075581+XinlongCS@users.noreply.github.com>
Date: Thu, 3 Jul 2025 12:55:31 -0500
Subject: [PATCH] Additional try except test case and fix

---
 .../templates/Exception.py                       |  4 ++--
 test/TryExcept.py                                | 16 ++++++++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/pylingual/control_flow_reconstruction/templates/Exception.py b/pylingual/control_flow_reconstruction/templates/Exception.py
index 9ebf04a..715c1d0 100644
--- a/pylingual/control_flow_reconstruction/templates/Exception.py
+++ b/pylingual/control_flow_reconstruction/templates/Exception.py
@@ -631,8 +631,8 @@ class TryElse3_6(ControlFlowTemplate):
         try_header=~N("try_body").with_cond(exact_instructions("SETUP_EXCEPT"), exact_instructions("SETUP_FINALLY")),
         try_body=N("try_footer.", None, "except_body"),
         try_footer=~N("else_body").with_in_deg(1),
-        except_body=~N("tail").with_in_deg(1).of_subtemplate(Except3_6),
-        else_body=~N("tail").with_in_deg(1),
+        except_body=~N("tail.").with_in_deg(1).of_subtemplate(Except3_6),
+        else_body=~N("tail.").with_in_deg(1),
         tail=N.tail(),
     )
 
diff --git a/test/TryExcept.py b/test/TryExcept.py
index 173895c..94cd00e 100644
--- a/test/TryExcept.py
+++ b/test/TryExcept.py
@@ -407,6 +407,14 @@ def z_TryExceptReturn():
     except:
         print(2)
 
+def z1_TryExceptReturn():
+    try:
+        print(1)
+        return 2
+    except:
+        print(2)
+        return 3
+
 def aa_TryExceptReturnMulti():
     try:
         print(1)
@@ -423,6 +431,14 @@ def ab_TryExceptReturnNamed():
     except A as a:
         print(2)
 
+def ab1_TryExceptReturnNamed():
+    try:
+        print(1)
+        return 2
+    except A as a:
+        print(2)
+        return 3
+
 def ac_TryFinallyBareFallthrough():
     try:
         print(1)