moving to scripts
This commit is contained in:
eneller
386
asq-env/lib/python3.9/site-packages/trio/_channel.py
Normal file
386
asq-env/lib/python3.9/site-packages/trio/_channel.py
Normal file
@@ -0,0 +1,386 @@
|
||||
from collections import deque, OrderedDict
|
||||
from math import inf
|
||||
|
||||
import attr
|
||||
from outcome import Error, Value
|
||||
|
||||
from .abc import SendChannel, ReceiveChannel, Channel
|
||||
from ._util import generic_function, NoPublicConstructor
|
||||
|
||||
import trio
|
||||
from ._core import enable_ki_protection
|
||||
|
||||
|
||||
@generic_function
|
||||
def open_memory_channel(max_buffer_size):
|
||||
"""Open a channel for passing objects between tasks within a process.
|
||||
|
||||
Memory channels are lightweight, cheap to allocate, and entirely
|
||||
in-memory. They don't involve any operating-system resources, or any kind
|
||||
of serialization. They just pass Python objects directly between tasks
|
||||
(with a possible stop in an internal buffer along the way).
|
||||
|
||||
Channel objects can be closed by calling `~trio.abc.AsyncResource.aclose`
|
||||
or using ``async with``. They are *not* automatically closed when garbage
|
||||
collected. Closing memory channels isn't mandatory, but it is generally a
|
||||
good idea, because it helps avoid situations where tasks get stuck waiting
|
||||
on a channel when there's no-one on the other side. See
|
||||
:ref:`channel-shutdown` for details.
|
||||
|
||||
Memory channel operations are all atomic with respect to
|
||||
cancellation, either `~trio.abc.ReceiveChannel.receive` will
|
||||
successfully return an object, or it will raise :exc:`Cancelled`
|
||||
while leaving the channel unchanged.
|
||||
|
||||
Args:
|
||||
max_buffer_size (int or math.inf): The maximum number of items that can
|
||||
be buffered in the channel before :meth:`~trio.abc.SendChannel.send`
|
||||
blocks. Choosing a sensible value here is important to ensure that
|
||||
backpressure is communicated promptly and avoid unnecessary latency;
|
||||
see :ref:`channel-buffering` for more details. If in doubt, use 0.
|
||||
|
||||
Returns:
|
||||
A pair ``(send_channel, receive_channel)``. If you have
|
||||
trouble remembering which order these go in, remember: data
|
||||
flows from left → right.
|
||||
|
||||
In addition to the standard channel methods, all memory channel objects
|
||||
provide a ``statistics()`` method, which returns an object with the
|
||||
following fields:
|
||||
|
||||
* ``current_buffer_used``: The number of items currently stored in the
|
||||
channel buffer.
|
||||
* ``max_buffer_size``: The maximum number of items allowed in the buffer,
|
||||
as passed to :func:`open_memory_channel`.
|
||||
* ``open_send_channels``: The number of open
|
||||
:class:`MemorySendChannel` endpoints pointing to this channel.
|
||||
Initially 1, but can be increased by
|
||||
:meth:`MemorySendChannel.clone`.
|
||||
* ``open_receive_channels``: Likewise, but for open
|
||||
:class:`MemoryReceiveChannel` endpoints.
|
||||
* ``tasks_waiting_send``: The number of tasks blocked in ``send`` on this
|
||||
channel (summing over all clones).
|
||||
* ``tasks_waiting_receive``: The number of tasks blocked in ``receive`` on
|
||||
this channel (summing over all clones).
|
||||
|
||||
"""
|
||||
if max_buffer_size != inf and not isinstance(max_buffer_size, int):
|
||||
raise TypeError("max_buffer_size must be an integer or math.inf")
|
||||
if max_buffer_size < 0:
|
||||
raise ValueError("max_buffer_size must be >= 0")
|
||||
state = MemoryChannelState(max_buffer_size)
|
||||
return (
|
||||
MemorySendChannel._create(state),
|
||||
MemoryReceiveChannel._create(state),
|
||||
)
|
||||
|
||||
|
||||
@attr.s(frozen=True, slots=True)
|
||||
class MemoryChannelStats:
|
||||
current_buffer_used = attr.ib()
|
||||
max_buffer_size = attr.ib()
|
||||
open_send_channels = attr.ib()
|
||||
open_receive_channels = attr.ib()
|
||||
tasks_waiting_send = attr.ib()
|
||||
tasks_waiting_receive = attr.ib()
|
||||
|
||||
|
||||
@attr.s(slots=True)
|
||||
class MemoryChannelState:
|
||||
max_buffer_size = attr.ib()
|
||||
data = attr.ib(factory=deque)
|
||||
# Counts of open endpoints using this state
|
||||
open_send_channels = attr.ib(default=0)
|
||||
open_receive_channels = attr.ib(default=0)
|
||||
# {task: value}
|
||||
send_tasks = attr.ib(factory=OrderedDict)
|
||||
# {task: None}
|
||||
receive_tasks = attr.ib(factory=OrderedDict)
|
||||
|
||||
def statistics(self):
|
||||
return MemoryChannelStats(
|
||||
current_buffer_used=len(self.data),
|
||||
max_buffer_size=self.max_buffer_size,
|
||||
open_send_channels=self.open_send_channels,
|
||||
open_receive_channels=self.open_receive_channels,
|
||||
tasks_waiting_send=len(self.send_tasks),
|
||||
tasks_waiting_receive=len(self.receive_tasks),
|
||||
)
|
||||
|
||||
|
||||
@attr.s(eq=False, repr=False)
|
||||
class MemorySendChannel(SendChannel, metaclass=NoPublicConstructor):
|
||||
_state = attr.ib()
|
||||
_closed = attr.ib(default=False)
|
||||
# This is just the tasks waiting on *this* object. As compared to
|
||||
# self._state.send_tasks, which includes tasks from this object and
|
||||
# all clones.
|
||||
_tasks = attr.ib(factory=set)
|
||||
|
||||
def __attrs_post_init__(self):
|
||||
self._state.open_send_channels += 1
|
||||
|
||||
def __repr__(self):
|
||||
return "<send channel at {:#x}, using buffer at {:#x}>".format(
|
||||
id(self), id(self._state)
|
||||
)
|
||||
|
||||
def statistics(self):
|
||||
# XX should we also report statistics specific to this object?
|
||||
return self._state.statistics()
|
||||
|
||||
@enable_ki_protection
|
||||
def send_nowait(self, value):
|
||||
"""Like `~trio.abc.SendChannel.send`, but if the channel's buffer is
|
||||
full, raises `WouldBlock` instead of blocking.
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
raise trio.ClosedResourceError
|
||||
if self._state.open_receive_channels == 0:
|
||||
raise trio.BrokenResourceError
|
||||
if self._state.receive_tasks:
|
||||
assert not self._state.data
|
||||
task, _ = self._state.receive_tasks.popitem(last=False)
|
||||
task.custom_sleep_data._tasks.remove(task)
|
||||
trio.lowlevel.reschedule(task, Value(value))
|
||||
elif len(self._state.data) < self._state.max_buffer_size:
|
||||
self._state.data.append(value)
|
||||
else:
|
||||
raise trio.WouldBlock
|
||||
|
||||
@enable_ki_protection
|
||||
async def send(self, value):
|
||||
"""See `SendChannel.send <trio.abc.SendChannel.send>`.
|
||||
|
||||
Memory channels allow multiple tasks to call `send` at the same time.
|
||||
|
||||
"""
|
||||
await trio.lowlevel.checkpoint_if_cancelled()
|
||||
try:
|
||||
self.send_nowait(value)
|
||||
except trio.WouldBlock:
|
||||
pass
|
||||
else:
|
||||
await trio.lowlevel.cancel_shielded_checkpoint()
|
||||
return
|
||||
|
||||
task = trio.lowlevel.current_task()
|
||||
self._tasks.add(task)
|
||||
self._state.send_tasks[task] = value
|
||||
task.custom_sleep_data = self
|
||||
|
||||
def abort_fn(_):
|
||||
self._tasks.remove(task)
|
||||
del self._state.send_tasks[task]
|
||||
return trio.lowlevel.Abort.SUCCEEDED
|
||||
|
||||
await trio.lowlevel.wait_task_rescheduled(abort_fn)
|
||||
|
||||
@enable_ki_protection
|
||||
def clone(self):
|
||||
"""Clone this send channel object.
|
||||
|
||||
This returns a new `MemorySendChannel` object, which acts as a
|
||||
duplicate of the original: sending on the new object does exactly the
|
||||
same thing as sending on the old object. (If you're familiar with
|
||||
`os.dup`, then this is a similar idea.)
|
||||
|
||||
However, closing one of the objects does not close the other, and
|
||||
receivers don't get `EndOfChannel` until *all* clones have been
|
||||
closed.
|
||||
|
||||
This is useful for communication patterns that involve multiple
|
||||
producers all sending objects to the same destination. If you give
|
||||
each producer its own clone of the `MemorySendChannel`, and then make
|
||||
sure to close each `MemorySendChannel` when it's finished, receivers
|
||||
will automatically get notified when all producers are finished. See
|
||||
:ref:`channel-mpmc` for examples.
|
||||
|
||||
Raises:
|
||||
trio.ClosedResourceError: if you already closed this
|
||||
`MemorySendChannel` object.
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
raise trio.ClosedResourceError
|
||||
return MemorySendChannel._create(self._state)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.close()
|
||||
|
||||
@enable_ki_protection
|
||||
def close(self):
|
||||
"""Close this send channel object synchronously.
|
||||
|
||||
All channel objects have an asynchronous `~.AsyncResource.aclose` method.
|
||||
Memory channels can also be closed synchronously. This has the same
|
||||
effect on the channel and other tasks using it, but `close` is not a
|
||||
trio checkpoint. This simplifies cleaning up in cancelled tasks.
|
||||
|
||||
Using ``with send_channel:`` will close the channel object on leaving
|
||||
the with block.
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
for task in self._tasks:
|
||||
trio.lowlevel.reschedule(task, Error(trio.ClosedResourceError()))
|
||||
del self._state.send_tasks[task]
|
||||
self._tasks.clear()
|
||||
self._state.open_send_channels -= 1
|
||||
if self._state.open_send_channels == 0:
|
||||
assert not self._state.send_tasks
|
||||
for task in self._state.receive_tasks:
|
||||
task.custom_sleep_data._tasks.remove(task)
|
||||
trio.lowlevel.reschedule(task, Error(trio.EndOfChannel()))
|
||||
self._state.receive_tasks.clear()
|
||||
|
||||
@enable_ki_protection
|
||||
async def aclose(self):
|
||||
self.close()
|
||||
await trio.lowlevel.checkpoint()
|
||||
|
||||
|
||||
@attr.s(eq=False, repr=False)
|
||||
class MemoryReceiveChannel(ReceiveChannel, metaclass=NoPublicConstructor):
|
||||
_state = attr.ib()
|
||||
_closed = attr.ib(default=False)
|
||||
_tasks = attr.ib(factory=set)
|
||||
|
||||
def __attrs_post_init__(self):
|
||||
self._state.open_receive_channels += 1
|
||||
|
||||
def statistics(self):
|
||||
return self._state.statistics()
|
||||
|
||||
def __repr__(self):
|
||||
return "<receive channel at {:#x}, using buffer at {:#x}>".format(
|
||||
id(self), id(self._state)
|
||||
)
|
||||
|
||||
@enable_ki_protection
|
||||
def receive_nowait(self):
|
||||
"""Like `~trio.abc.ReceiveChannel.receive`, but if there's nothing
|
||||
ready to receive, raises `WouldBlock` instead of blocking.
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
raise trio.ClosedResourceError
|
||||
if self._state.send_tasks:
|
||||
task, value = self._state.send_tasks.popitem(last=False)
|
||||
task.custom_sleep_data._tasks.remove(task)
|
||||
trio.lowlevel.reschedule(task)
|
||||
self._state.data.append(value)
|
||||
# Fall through
|
||||
if self._state.data:
|
||||
return self._state.data.popleft()
|
||||
if not self._state.open_send_channels:
|
||||
raise trio.EndOfChannel
|
||||
raise trio.WouldBlock
|
||||
|
||||
@enable_ki_protection
|
||||
async def receive(self):
|
||||
"""See `ReceiveChannel.receive <trio.abc.ReceiveChannel.receive>`.
|
||||
|
||||
Memory channels allow multiple tasks to call `receive` at the same
|
||||
time. The first task will get the first item sent, the second task
|
||||
will get the second item sent, and so on.
|
||||
|
||||
"""
|
||||
await trio.lowlevel.checkpoint_if_cancelled()
|
||||
try:
|
||||
value = self.receive_nowait()
|
||||
except trio.WouldBlock:
|
||||
pass
|
||||
else:
|
||||
await trio.lowlevel.cancel_shielded_checkpoint()
|
||||
return value
|
||||
|
||||
task = trio.lowlevel.current_task()
|
||||
self._tasks.add(task)
|
||||
self._state.receive_tasks[task] = None
|
||||
task.custom_sleep_data = self
|
||||
|
||||
def abort_fn(_):
|
||||
self._tasks.remove(task)
|
||||
del self._state.receive_tasks[task]
|
||||
return trio.lowlevel.Abort.SUCCEEDED
|
||||
|
||||
return await trio.lowlevel.wait_task_rescheduled(abort_fn)
|
||||
|
||||
@enable_ki_protection
|
||||
def clone(self):
|
||||
"""Clone this receive channel object.
|
||||
|
||||
This returns a new `MemoryReceiveChannel` object, which acts as a
|
||||
duplicate of the original: receiving on the new object does exactly
|
||||
the same thing as receiving on the old object.
|
||||
|
||||
However, closing one of the objects does not close the other, and the
|
||||
underlying channel is not closed until all clones are closed. (If
|
||||
you're familiar with `os.dup`, then this is a similar idea.)
|
||||
|
||||
This is useful for communication patterns that involve multiple
|
||||
consumers all receiving objects from the same underlying channel. See
|
||||
:ref:`channel-mpmc` for examples.
|
||||
|
||||
.. warning:: The clones all share the same underlying channel.
|
||||
Whenever a clone :meth:`receive`\\s a value, it is removed from the
|
||||
channel and the other clones do *not* receive that value. If you
|
||||
want to send multiple copies of the same stream of values to
|
||||
multiple destinations, like :func:`itertools.tee`, then you need to
|
||||
find some other solution; this method does *not* do that.
|
||||
|
||||
Raises:
|
||||
trio.ClosedResourceError: if you already closed this
|
||||
`MemoryReceiveChannel` object.
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
raise trio.ClosedResourceError
|
||||
return MemoryReceiveChannel._create(self._state)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.close()
|
||||
|
||||
@enable_ki_protection
|
||||
def close(self):
|
||||
"""Close this receive channel object synchronously.
|
||||
|
||||
All channel objects have an asynchronous `~.AsyncResource.aclose` method.
|
||||
Memory channels can also be closed synchronously. This has the same
|
||||
effect on the channel and other tasks using it, but `close` is not a
|
||||
trio checkpoint. This simplifies cleaning up in cancelled tasks.
|
||||
|
||||
Using ``with receive_channel:`` will close the channel object on
|
||||
leaving the with block.
|
||||
|
||||
"""
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
for task in self._tasks:
|
||||
trio.lowlevel.reschedule(task, Error(trio.ClosedResourceError()))
|
||||
del self._state.receive_tasks[task]
|
||||
self._tasks.clear()
|
||||
self._state.open_receive_channels -= 1
|
||||
if self._state.open_receive_channels == 0:
|
||||
assert not self._state.receive_tasks
|
||||
for task in self._state.send_tasks:
|
||||
task.custom_sleep_data._tasks.remove(task)
|
||||
trio.lowlevel.reschedule(task, Error(trio.BrokenResourceError()))
|
||||
self._state.send_tasks.clear()
|
||||
self._state.data.clear()
|
||||
|
||||
@enable_ki_protection
|
||||
async def aclose(self):
|
||||
self.close()
|
||||
await trio.lowlevel.checkpoint()
|
||||
Reference in New Issue
Block a user