Skip to content

Commit 96fed66

Browse files
gh-110378: Close invalid generators in contextmanager and asynccontextmanager (GH-110499)
contextmanager and asynccontextmanager context managers now close an invalid underlying generator object that yields more then one value.
1 parent def7ea5 commit 96fed66

File tree

4 files changed

+43
-7
lines changed

4 files changed

+43
-7
lines changed

Lib/contextlib.py

+16-4
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,10 @@ def __exit__(self, typ, value, traceback):
149149
except StopIteration:
150150
return False
151151
else:
152-
raise RuntimeError("generator didn't stop")
152+
try:
153+
raise RuntimeError("generator didn't stop")
154+
finally:
155+
self.gen.close()
153156
else:
154157
if value is None:
155158
# Need to force instantiation so we can reliably
@@ -191,7 +194,10 @@ def __exit__(self, typ, value, traceback):
191194
raise
192195
exc.__traceback__ = traceback
193196
return False
194-
raise RuntimeError("generator didn't stop after throw()")
197+
try:
198+
raise RuntimeError("generator didn't stop after throw()")
199+
finally:
200+
self.gen.close()
195201

196202
class _AsyncGeneratorContextManager(
197203
_GeneratorContextManagerBase,
@@ -216,7 +222,10 @@ async def __aexit__(self, typ, value, traceback):
216222
except StopAsyncIteration:
217223
return False
218224
else:
219-
raise RuntimeError("generator didn't stop")
225+
try:
226+
raise RuntimeError("generator didn't stop")
227+
finally:
228+
await self.gen.aclose()
220229
else:
221230
if value is None:
222231
# Need to force instantiation so we can reliably
@@ -258,7 +267,10 @@ async def __aexit__(self, typ, value, traceback):
258267
raise
259268
exc.__traceback__ = traceback
260269
return False
261-
raise RuntimeError("generator didn't stop after athrow()")
270+
try:
271+
raise RuntimeError("generator didn't stop after athrow()")
272+
finally:
273+
await self.gen.aclose()
262274

263275

264276
def contextmanager(func):

Lib/test/test_contextlib.py

+18-3
Original file line numberDiff line numberDiff line change
@@ -167,9 +167,24 @@ def whoo():
167167
yield
168168
ctx = whoo()
169169
ctx.__enter__()
170-
self.assertRaises(
171-
RuntimeError, ctx.__exit__, TypeError, TypeError("foo"), None
172-
)
170+
with self.assertRaises(RuntimeError):
171+
ctx.__exit__(TypeError, TypeError("foo"), None)
172+
if support.check_impl_detail(cpython=True):
173+
# The "gen" attribute is an implementation detail.
174+
self.assertFalse(ctx.gen.gi_suspended)
175+
176+
def test_contextmanager_trap_second_yield(self):
177+
@contextmanager
178+
def whoo():
179+
yield
180+
yield
181+
ctx = whoo()
182+
ctx.__enter__()
183+
with self.assertRaises(RuntimeError):
184+
ctx.__exit__(None, None, None)
185+
if support.check_impl_detail(cpython=True):
186+
# The "gen" attribute is an implementation detail.
187+
self.assertFalse(ctx.gen.gi_suspended)
173188

174189
def test_contextmanager_except(self):
175190
state = []

Lib/test/test_contextlib_async.py

+6
Original file line numberDiff line numberDiff line change
@@ -199,6 +199,9 @@ async def whoo():
199199
await ctx.__aenter__()
200200
with self.assertRaises(RuntimeError):
201201
await ctx.__aexit__(TypeError, TypeError('foo'), None)
202+
if support.check_impl_detail(cpython=True):
203+
# The "gen" attribute is an implementation detail.
204+
self.assertFalse(ctx.gen.ag_suspended)
202205

203206
async def test_contextmanager_trap_no_yield(self):
204207
@asynccontextmanager
@@ -218,6 +221,9 @@ async def whoo():
218221
await ctx.__aenter__()
219222
with self.assertRaises(RuntimeError):
220223
await ctx.__aexit__(None, None, None)
224+
if support.check_impl_detail(cpython=True):
225+
# The "gen" attribute is an implementation detail.
226+
self.assertFalse(ctx.gen.ag_suspended)
221227

222228
async def test_contextmanager_non_normalised(self):
223229
@asynccontextmanager
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
:func:`~contextlib.contextmanager` and
2+
:func:`~contextlib.asynccontextmanager` context managers now close an invalid
3+
underlying generator object that yields more then one value.

0 commit comments

Comments
 (0)