Skip to content

Commit

Permalink
[3.12] gh-110378: Close invalid generators in contextmanager and asyn…
Browse files Browse the repository at this point in the history
…ccontextmanager (GH-110499) (#110588)

contextmanager and asynccontextmanager context managers now close an invalid
underlying generator object that yields more then one value.
(cherry picked from commit 96fed66)

Co-authored-by: Serhiy Storchaka <storchaka@gmail.com>
  • Loading branch information
miss-islington and serhiy-storchaka committed Oct 10, 2023
1 parent 3688672 commit 2fc8081
Show file tree
Hide file tree
Showing 4 changed files with 43 additions and 7 deletions.
20 changes: 16 additions & 4 deletions Lib/contextlib.py
Expand Up @@ -145,7 +145,10 @@ def __exit__(self, typ, value, traceback):
except StopIteration:
return False
else:
raise RuntimeError("generator didn't stop")
try:
raise RuntimeError("generator didn't stop")
finally:
self.gen.close()
else:
if value is None:
# Need to force instantiation so we can reliably
Expand Down Expand Up @@ -187,7 +190,10 @@ def __exit__(self, typ, value, traceback):
raise
exc.__traceback__ = traceback
return False
raise RuntimeError("generator didn't stop after throw()")
try:
raise RuntimeError("generator didn't stop after throw()")
finally:
self.gen.close()

class _AsyncGeneratorContextManager(
_GeneratorContextManagerBase,
Expand All @@ -212,7 +218,10 @@ async def __aexit__(self, typ, value, traceback):
except StopAsyncIteration:
return False
else:
raise RuntimeError("generator didn't stop")
try:
raise RuntimeError("generator didn't stop")
finally:
await self.gen.aclose()
else:
if value is None:
# Need to force instantiation so we can reliably
Expand Down Expand Up @@ -254,7 +263,10 @@ async def __aexit__(self, typ, value, traceback):
raise
exc.__traceback__ = traceback
return False
raise RuntimeError("generator didn't stop after athrow()")
try:
raise RuntimeError("generator didn't stop after athrow()")
finally:
await self.gen.aclose()


def contextmanager(func):
Expand Down
21 changes: 18 additions & 3 deletions Lib/test/test_contextlib.py
Expand Up @@ -157,9 +157,24 @@ def whoo():
yield
ctx = whoo()
ctx.__enter__()
self.assertRaises(
RuntimeError, ctx.__exit__, TypeError, TypeError("foo"), None
)
with self.assertRaises(RuntimeError):
ctx.__exit__(TypeError, TypeError("foo"), None)
if support.check_impl_detail(cpython=True):
# The "gen" attribute is an implementation detail.
self.assertFalse(ctx.gen.gi_suspended)

def test_contextmanager_trap_second_yield(self):
@contextmanager
def whoo():
yield
yield
ctx = whoo()
ctx.__enter__()
with self.assertRaises(RuntimeError):
ctx.__exit__(None, None, None)
if support.check_impl_detail(cpython=True):
# The "gen" attribute is an implementation detail.
self.assertFalse(ctx.gen.gi_suspended)

def test_contextmanager_except(self):
state = []
Expand Down
6 changes: 6 additions & 0 deletions Lib/test/test_contextlib_async.py
Expand Up @@ -204,6 +204,9 @@ async def whoo():
await ctx.__aenter__()
with self.assertRaises(RuntimeError):
await ctx.__aexit__(TypeError, TypeError('foo'), None)
if support.check_impl_detail(cpython=True):
# The "gen" attribute is an implementation detail.
self.assertFalse(ctx.gen.ag_suspended)

@_async_test
async def test_contextmanager_trap_no_yield(self):
Expand All @@ -225,6 +228,9 @@ async def whoo():
await ctx.__aenter__()
with self.assertRaises(RuntimeError):
await ctx.__aexit__(None, None, None)
if support.check_impl_detail(cpython=True):
# The "gen" attribute is an implementation detail.
self.assertFalse(ctx.gen.ag_suspended)

@_async_test
async def test_contextmanager_non_normalised(self):
Expand Down
@@ -0,0 +1,3 @@
:func:`~contextlib.contextmanager` and
:func:`~contextlib.asynccontextmanager` context managers now close an invalid
underlying generator object that yields more then one value.

0 comments on commit 2fc8081

Please sign in to comment.