From baaadbf70d024142ec6a519006e960ec74309f92 Mon Sep 17 00:00:00 2001
From: Nick Coghlan <ncoghlan@gmail.com>
Date: Wed, 13 May 2015 15:54:02 +1000
Subject: [PATCH] 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
---
 Lib/test/test_coroutines.py | 120 +++++++++++++++++++++++++++++++++++-
 Python/ceval.c              |   2 +
 2 files changed, 120 insertions(+), 2 deletions(-)

diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
index aa2a5e8ef57..6a6f868c129 100644
--- a/Lib/test/test_coroutines.py
+++ b/Lib/test/test_coroutines.py
@@ -497,17 +497,133 @@ async def __aenter__(self):
                 return self
 
             def __aexit__(self, *e):
-                return 456
+                return 444
 
         async def foo():
             async with CM():
-                pass
+                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():
+                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
index 77085c25a03..afb0f89aa8d 100644
--- a/Python/ceval.c
+++ b/Python/ceval.c
@@ -3156,6 +3156,7 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
             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 @@ PyEval_EvalFrameEx(PyFrameObject *f, int throwflag)
                 err = 0;
 
             Py_DECREF(res);
+            Py_DECREF(exc);
 
             if (err < 0)
                 goto error;