You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
108 lines
3.8 KiB
Python
108 lines
3.8 KiB
Python
import attr
|
|
|
|
import trio
|
|
from .abc import HalfCloseableStream
|
|
|
|
from trio._util import Final
|
|
|
|
|
|
async def aclose_forcefully(resource):
|
|
"""Close an async resource or async generator immediately, without
|
|
blocking to do any graceful cleanup.
|
|
|
|
:class:`~trio.abc.AsyncResource` objects guarantee that if their
|
|
:meth:`~trio.abc.AsyncResource.aclose` method is cancelled, then they will
|
|
still close the resource (albeit in a potentially ungraceful
|
|
fashion). :func:`aclose_forcefully` is a convenience function that
|
|
exploits this behavior to let you force a resource to be closed without
|
|
blocking: it works by calling ``await resource.aclose()`` and then
|
|
cancelling it immediately.
|
|
|
|
Most users won't need this, but it may be useful on cleanup paths where
|
|
you can't afford to block, or if you want to close a resource and don't
|
|
care about handling it gracefully. For example, if
|
|
:class:`~trio.SSLStream` encounters an error and cannot perform its
|
|
own graceful close, then there's no point in waiting to gracefully shut
|
|
down the underlying transport either, so it calls ``await
|
|
aclose_forcefully(self.transport_stream)``.
|
|
|
|
Note that this function is async, and that it acts as a checkpoint, but
|
|
unlike most async functions it cannot block indefinitely (at least,
|
|
assuming the underlying resource object is correctly implemented).
|
|
|
|
"""
|
|
with trio.CancelScope() as cs:
|
|
cs.cancel()
|
|
await resource.aclose()
|
|
|
|
|
|
@attr.s(eq=False, hash=False)
|
|
class StapledStream(HalfCloseableStream, metaclass=Final):
|
|
"""This class `staples <https://en.wikipedia.org/wiki/Staple_(fastener)>`__
|
|
together two unidirectional streams to make single bidirectional stream.
|
|
|
|
Args:
|
|
send_stream (~trio.abc.SendStream): The stream to use for sending.
|
|
receive_stream (~trio.abc.ReceiveStream): The stream to use for
|
|
receiving.
|
|
|
|
Example:
|
|
|
|
A silly way to make a stream that echoes back whatever you write to
|
|
it::
|
|
|
|
left, right = trio.testing.memory_stream_pair()
|
|
echo_stream = StapledStream(SocketStream(left), SocketStream(right))
|
|
await echo_stream.send_all(b"x")
|
|
assert await echo_stream.receive_some() == b"x"
|
|
|
|
:class:`StapledStream` objects implement the methods in the
|
|
:class:`~trio.abc.HalfCloseableStream` interface. They also have two
|
|
additional public attributes:
|
|
|
|
.. attribute:: send_stream
|
|
|
|
The underlying :class:`~trio.abc.SendStream`. :meth:`send_all` and
|
|
:meth:`wait_send_all_might_not_block` are delegated to this object.
|
|
|
|
.. attribute:: receive_stream
|
|
|
|
The underlying :class:`~trio.abc.ReceiveStream`. :meth:`receive_some`
|
|
is delegated to this object.
|
|
|
|
"""
|
|
|
|
send_stream = attr.ib()
|
|
receive_stream = attr.ib()
|
|
|
|
async def send_all(self, data):
|
|
"""Calls ``self.send_stream.send_all``."""
|
|
return await self.send_stream.send_all(data)
|
|
|
|
async def wait_send_all_might_not_block(self):
|
|
"""Calls ``self.send_stream.wait_send_all_might_not_block``."""
|
|
return await self.send_stream.wait_send_all_might_not_block()
|
|
|
|
async def send_eof(self):
|
|
"""Shuts down the send side of the stream.
|
|
|
|
If ``self.send_stream.send_eof`` exists, then calls it. Otherwise,
|
|
calls ``self.send_stream.aclose()``.
|
|
|
|
"""
|
|
if hasattr(self.send_stream, "send_eof"):
|
|
return await self.send_stream.send_eof()
|
|
else:
|
|
return await self.send_stream.aclose()
|
|
|
|
async def receive_some(self, max_bytes=None):
|
|
"""Calls ``self.receive_stream.receive_some``."""
|
|
return await self.receive_stream.receive_some(max_bytes)
|
|
|
|
async def aclose(self):
|
|
"""Calls ``aclose`` on both underlying streams."""
|
|
try:
|
|
await self.send_stream.aclose()
|
|
finally:
|
|
await self.receive_stream.aclose()
|