mirrors/cpython
1
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2025-07-24 11:44:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

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

Browse source
This commit is contained in:
Andrew Svetlov 2018-01-25 00:30:30 +02:00 • committed by GitHub
parent 04af5b1ba9
commit fe133aad52
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 73 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

36
Lib/asyncio/streams.py
View file

@ -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()