Skip to content

gh-110378: Close invalid generators in contextmanager and asynccontextmanager #110499

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 16 additions & 4 deletions Lib/contextlib.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,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()
Comment on lines +152 to +155
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍🏻

I sometimes use the try:finally: construct around flow breaking instructions as well. Here it nicely chains exceptions if self.gen.close() raises.

else:
if value is None:
# Need to force instantiation so we can reliably
Expand Down Expand Up @@ -191,7 +194,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 @@ -216,7 +222,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 @@ -258,7 +267,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
Original file line number Diff line number Diff line change
Expand Up @@ -167,9 +167,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
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,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 @@ -237,6 +240,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
Original file line number Diff line number Diff line change
@@ -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.