bpo-32391: Implement StreamWriter.wait_closed() (#5281) · python/cpython@fe133aa (original) (raw)
`@@ -224,6 +224,7 @@ def init(self, stream_reader, client_connected_cb=None, loop=None):
`
224
224
`self._stream_writer = None
`
225
225
`self._client_connected_cb = client_connected_cb
`
226
226
`self._over_ssl = False
`
``
227
`+
self._closed = self._loop.create_future()
`
227
228
``
228
229
`def connection_made(self, transport):
`
229
230
`self._stream_reader.set_transport(transport)
`
`@@ -243,6 +244,11 @@ def connection_lost(self, exc):
`
243
244
`self._stream_reader.feed_eof()
`
244
245
`else:
`
245
246
`self._stream_reader.set_exception(exc)
`
``
247
`+
if not self._closed.done():
`
``
248
`+
if exc is None:
`
``
249
`+
self._closed.set_result(None)
`
``
250
`+
else:
`
``
251
`+
self._closed.set_exception(exc)
`
246
252
`super().connection_lost(exc)
`
247
253
`self._stream_reader = None
`
248
254
`self._stream_writer = None
`
`@@ -259,6 +265,13 @@ def eof_received(self):
`
259
265
`return False
`
260
266
`return True
`
261
267
``
``
268
`+
def del(self):
`
``
269
`+
Prevent reports about unhandled exceptions.
`
``
270
`+
Better than self._closed._log_traceback = False hack
`
``
271
`+
closed = self._closed
`
``
272
`+
if closed.done() and not closed.cancelled():
`
``
273
`+
closed.exception()
`
``
274
+
262
275
``
263
276
`class StreamWriter:
`
264
277
`"""Wraps a Transport.
`
`@@ -303,6 +316,12 @@ def can_write_eof(self):
`
303
316
`def close(self):
`
304
317
`return self._transport.close()
`
305
318
``
``
319
`+
def is_closing(self):
`
``
320
`+
return self._transport.is_closing()
`
``
321
+
``
322
`+
async def wait_closed(self):
`
``
323
`+
await self._protocol._closed
`
``
324
+
306
325
`def get_extra_info(self, name, default=None):
`
307
326
`return self._transport.get_extra_info(name, default)
`
308
327
``
`@@ -318,15 +337,14 @@ async def drain(self):
`
318
337
`exc = self._reader.exception()
`
319
338
`if exc is not None:
`
320
339
`raise exc
`
321
``
`-
if self._transport is not None:
`
322
``
`-
if self._transport.is_closing():
`
323
``
`-
Yield to the event loop so connection_lost() may be
`
324
``
`-
called. Without this, _drain_helper() would return
`
325
``
`-
immediately, and code that calls
`
326
``
`-
write(...); await drain()
`
327
``
`-
in a loop would never call connection_lost(), so it
`
328
``
`-
would not see an error when the socket is closed.
`
329
``
`-
await sleep(0, loop=self._loop)
`
``
340
`+
if self._transport.is_closing():
`
``
341
`+
Yield to the event loop so connection_lost() may be
`
``
342
`+
called. Without this, _drain_helper() would return
`
``
343
`+
immediately, and code that calls
`
``
344
`+
write(...); await drain()
`
``
345
`+
in a loop would never call connection_lost(), so it
`
``
346
`+
would not see an error when the socket is closed.
`
``
347
`+
await sleep(0, loop=self._loop)
`
330
348
`await self._protocol._drain_helper()
`
331
349
``
332
350
``