[Python-checkins] cpython: Issue 24017: fix for "async with" refcounting

[nick.coghlan]

https://hg.python.org/cpython/rev/e39fd5a8501a
changeset:   96004:e39fd5a8501a
user:        Nick Coghlan <ncoghlan at gmail.com>
date:        Wed May 13 15:54:02 2015 +1000
summary:
  Issue 24017: fix for "async with" refcounting

* adds missing INCREF in WITH_CLEANUP_START
* adds missing DECREF in WITH_CLEANUP_FINISH
* adds several new tests Yury created while investigating this

files:
  Lib/test/test_coroutines.py |  118 +++++++++++++++++++++++-
  Python/ceval.c              |    2 +
  2 files changed, 119 insertions(+), 1 deletions(-)


diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
--- a/Lib/test/test_coroutines.py
+++ b/Lib/test/test_coroutines.py
@@ -497,17 +497,133 @@
                 return self
 
             def __aexit__(self, *e):
+                return 444
+
+        async def foo():
+            async with CM():
+                1/0
+
+        try:
+            run_async(foo())
+        except TypeError as exc:
+            self.assertRegex(
+                exc.args[0], "object int can't be used in 'await' expression")
+            self.assertTrue(exc.__context__ is not None)
+            self.assertTrue(isinstance(exc.__context__, ZeroDivisionError))
+        else:
+            self.fail('invalid asynchronous context manager did not fail')
+
+
+    def test_with_8(self):
+        CNT = 0
+
+        class CM:
+            async def __aenter__(self):
+                return self
+
+            def __aexit__(self, *e):
                 return 456
 
         async def foo():
+            nonlocal CNT
             async with CM():
-                pass
+                CNT += 1
+
 
         with self.assertRaisesRegex(
             TypeError, "object int can't be used in 'await' expression"):
 
             run_async(foo())
 
+        self.assertEqual(CNT, 1)
+
+
+    def test_with_9(self):
+        CNT = 0
+
+        class CM:
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, *e):
+                1/0
+
+        async def foo():
+            nonlocal CNT
+            async with CM():
+                CNT += 1
+
+        with self.assertRaises(ZeroDivisionError):
+            run_async(foo())
+
+        self.assertEqual(CNT, 1)
+
+    def test_with_10(self):
+        CNT = 0
+
+        class CM:
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, *e):
+                1/0
+
+        async def foo():
+            nonlocal CNT
+            async with CM():
+                async with CM():
+                    raise RuntimeError
+
+        try:
+            run_async(foo())
+        except ZeroDivisionError as exc:
+            self.assertTrue(exc.__context__ is not None)
+            self.assertTrue(isinstance(exc.__context__, ZeroDivisionError))
+            self.assertTrue(isinstance(exc.__context__.__context__,
+                                       RuntimeError))
+        else:
+            self.fail('exception from __aexit__ did not propagate')
+
+    def test_with_11(self):
+        CNT = 0
+
+        class CM:
+            async def __aenter__(self):
+                raise NotImplementedError
+
+            async def __aexit__(self, *e):
+                1/0
+
+        async def foo():
+            nonlocal CNT
+            async with CM():
+                raise RuntimeError
+
+        try:
+            run_async(foo())
+        except NotImplementedError as exc:
+            self.assertTrue(exc.__context__ is None)
+        else:
+            self.fail('exception from __aenter__ did not propagate')
+
+    def test_with_12(self):
+        CNT = 0
+
+        class CM:
+            async def __aenter__(self):
+                return self
+
+            async def __aexit__(self, *e):
+                return True
+
+        async def foo():
+            nonlocal CNT
+            async with CM() as cm:
+                self.assertIs(cm.__class__, CM)
+                raise RuntimeError
+
+        run_async(foo())
+
     def test_for_1(self):
         aiter_calls = 0
 
diff --git a/Python/ceval.c b/Python/ceval.c
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -3156,6 +3156,7 @@
             if (res == NULL)
                 goto error;
 
+            Py_INCREF(exc); /* Duplicating the exception on the stack */
             PUSH(exc);
             PUSH(res);
             PREDICT(WITH_CLEANUP_FINISH);
@@ -3174,6 +3175,7 @@
                 err = 0;
 
             Py_DECREF(res);
+            Py_DECREF(exc);
 
             if (err < 0)
                 goto error;

-- 
Repository URL: https://hg.python.org/cpython