bpo-32391: Implement StreamWriter.wait_closed() (#5281)

1 files changed, 27 insertions, 9 deletions

diff --git a/Lib/asyncio/streams.py b/Lib/asyncio/streams.py
index eef2b89..9a53ee4 100644
--- a/Lib/asyncio/streams.py
+++ b/Lib/asyncio/streams.py
@@ -224,6 +224,7 @@ class StreamReaderProtocol(FlowControlMixin, protocols.Protocol):
         self._stream_writer = None
         self._client_connected_cb = client_connected_cb
         self._over_ssl = False
+        self._closed = self._loop.create_future()
 
     def connection_made(self, transport):
         self._stream_reader.set_transport(transport)
@@ -243,6 +244,11 @@ class StreamReaderProtocol(FlowControlMixin, protocols.Protocol):
                 self._stream_reader.feed_eof()
             else:
                 self._stream_reader.set_exception(exc)
+        if not self._closed.done():
+            if exc is None:
+                self._closed.set_result(None)
+            else:
+                self._closed.set_exception(exc)
         super().connection_lost(exc)
         self._stream_reader = None
         self._stream_writer = None
@@ -259,6 +265,13 @@ class StreamReaderProtocol(FlowControlMixin, protocols.Protocol):
             return False
         return True
 
+    def __del__(self):
+        # Prevent reports about unhandled exceptions.
+        # Better than self._closed._log_traceback = False hack
+        closed = self._closed
+        if closed.done() and not closed.cancelled():
+            closed.exception()
+
 
 class StreamWriter:
     """Wraps a Transport.
@@ -303,6 +316,12 @@ class StreamWriter:
     def close(self):
         return self._transport.close()
 
+    def is_closing(self):
+        return self._transport.is_closing()
+
+    async def wait_closed(self):
+        await self._protocol._closed
+
     def get_extra_info(self, name, default=None):
         return self._transport.get_extra_info(name, default)
 
@@ -318,15 +337,14 @@ class StreamWriter:
             exc = self._reader.exception()
             if exc is not None:
                 raise exc
-        if self._transport is not None:
-            if self._transport.is_closing():
-                # Yield to the event loop so connection_lost() may be
-                # called.  Without this, _drain_helper() would return
-                # immediately, and code that calls
-                #     write(...); await drain()
-                # in a loop would never call connection_lost(), so it
-                # would not see an error when the socket is closed.
-                await sleep(0, loop=self._loop)
+        if self._transport.is_closing():
+            # Yield to the event loop so connection_lost() may be
+            # called.  Without this, _drain_helper() would return
+            # immediately, and code that calls
+            #     write(...); await drain()
+            # in a loop would never call connection_lost(), so it
+            # would not see an error when the socket is closed.
+            await sleep(0, loop=self._loop)
         await self._protocol._drain_helper()