gh-108654: restore comprehension locals before handling exception (#108659)

Co-authored-by: Dong-hee Na <donghee.na92@gmail.com>

1 files changed, 35 insertions, 0 deletions

diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index 9f28ced..bedd99b 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -561,6 +561,41 @@ class ListComprehensionTest(unittest.TestCase):
             }
         )
 
+    def test_comp_in_try_except(self):
+        template = """
+            value = ["a"]
+            try:
+                [{func}(value) for value in value]
+            except:
+                pass
+        """
+        for func in ["str", "int"]:
+            code = template.format(func=func)
+            raises = func != "str"
+            with self.subTest(raises=raises):
+                self._check_in_scopes(code, {"value": ["a"]})
+
+    def test_comp_in_try_finally(self):
+        code = """
+            def f(value):
+                try:
+                    [{func}(value) for value in value]
+                finally:
+                    return value
+            ret = f(["a"])
+        """
+        self._check_in_scopes(code, {"ret": ["a"]})
+
+    def test_exception_in_post_comp_call(self):
+        code = """
+            value = [1, None]
+            try:
+                [v for v in value].sort()
+            except:
+                pass
+        """
+        self._check_in_scopes(code, {"value": [1, None]})
+
 
 __test__ = {'doctests' : doctests}