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
diff --git a/Lib/test/test_coroutines.py b/Lib/test/test_coroutines.py
index aa2a5e8..6a6f868 100644
--- a/Lib/test/test_coroutines.py
+++ b/Lib/test/test_coroutines.py
@@ -497,17 +497,133 @@
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 77085c2..afb0f89 100644
--- 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;