From 96fed66a65097eac2dc528ce29c9ba676bb07689 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Tue, 10 Oct 2023 10:43:04 +0300
Subject: [PATCH] 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.
---
 Lib/contextlib.py                             | 20 ++++++++++++++----
 Lib/test/test_contextlib.py                   | 21 ++++++++++++++++---
 Lib/test/test_contextlib_async.py             |  6 ++++++
 ...-10-07-13-50-12.gh-issue-110378.Y4L8fl.rst |  3 +++
 4 files changed, 43 insertions(+), 7 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Library/2023-10-07-13-50-12.gh-issue-110378.Y4L8fl.rst

diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index f82e7bca357..6994690ebf7 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -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()
         else:
             if value is None:
                 # Need to force instantiation so we can reliably
@@ -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,
@@ -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
@@ -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):
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index dbc7dfcc24b..5d94ec7cae4 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -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 = []
diff --git a/Lib/test/test_contextlib_async.py b/Lib/test/test_contextlib_async.py
index 8ccce5ae1e7..540964a9bfd 100644
--- a/Lib/test/test_contextlib_async.py
+++ b/Lib/test/test_contextlib_async.py
@@ -199,6 +199,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 def test_contextmanager_trap_no_yield(self):
         @asynccontextmanager
@@ -218,6 +221,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 def test_contextmanager_non_normalised(self):
         @asynccontextmanager
diff --git a/Misc/NEWS.d/next/Library/2023-10-07-13-50-12.gh-issue-110378.Y4L8fl.rst b/Misc/NEWS.d/next/Library/2023-10-07-13-50-12.gh-issue-110378.Y4L8fl.rst
new file mode 100644
index 00000000000..ef5395fc3c6
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2023-10-07-13-50-12.gh-issue-110378.Y4L8fl.rst
@@ -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.