poni
/
harvesting_the_net
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity
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.
master
Branches Tags
${ item.name }
Create tag ${ searchTerm }
Create branch ${ searchTerm }
from '0ee4301a8c'
${ noResults }
harvesting_the_net/axios-example/lib/python3.10/site-packages/trio/testing/_memory_streams.py

592 lines
21 KiB
Python
Raw Normal View History Unescape Escape

first commit
2 years ago
import operator
from .. import _core
from .._highlevel_generic import StapledStream
from .. import _util
from ..abc import SendStream, ReceiveStream
################################################################
# In-memory streams - Unbounded buffer version
################################################################
class _UnboundedByteQueue:
def __init__(self):
self._data = bytearray()
self._closed = False
self._lot = _core.ParkingLot()
self._fetch_lock = _util.ConflictDetector(
"another task is already fetching data"
)
# This object treats "close" as being like closing the send side of a
# channel: so after close(), calling put() raises ClosedResourceError, and
# calling the get() variants drains the buffer and then returns an empty
# bytearray.
def close(self):
self._closed = True
self._lot.unpark_all()
def close_and_wipe(self):
self._data = bytearray()
self.close()
def put(self, data):
if self._closed:
raise _core.ClosedResourceError("virtual connection closed")
self._data += data
self._lot.unpark_all()
def _check_max_bytes(self, max_bytes):
if max_bytes is None:
return
max_bytes = operator.index(max_bytes)
if max_bytes < 1:
raise ValueError("max_bytes must be >= 1")
def _get_impl(self, max_bytes):
assert self._closed or self._data
if max_bytes is None:
max_bytes = len(self._data)
if self._data:
chunk = self._data[:max_bytes]
del self._data[:max_bytes]
assert chunk
return chunk
else:
return bytearray()
def get_nowait(self, max_bytes=None):
with self._fetch_lock:
self._check_max_bytes(max_bytes)
if not self._closed and not self._data:
raise _core.WouldBlock
return self._get_impl(max_bytes)
async def get(self, max_bytes=None):
with self._fetch_lock:
self._check_max_bytes(max_bytes)
if not self._closed and not self._data:
await self._lot.park()
else:
await _core.checkpoint()
return self._get_impl(max_bytes)
class MemorySendStream(SendStream, metaclass=_util.Final):
"""An in-memory :class:`~trio.abc.SendStream`.
Args:
send_all_hook: An async function, or None. Called from
:meth:`send_all`. Can do whatever you like.
wait_send_all_might_not_block_hook: An async function, or None. Called
from :meth:`wait_send_all_might_not_block`. Can do whatever you
like.
close_hook: A synchronous function, or None. Called from :meth:`close`
and :meth:`aclose`. Can do whatever you like.
.. attribute:: send_all_hook
wait_send_all_might_not_block_hook
close_hook
All of these hooks are also exposed as attributes on the object, and
you can change them at any time.
"""
def __init__(
self,
send_all_hook=None,
wait_send_all_might_not_block_hook=None,
close_hook=None,
):
self._conflict_detector = _util.ConflictDetector(
"another task is using this stream"
)
self._outgoing = _UnboundedByteQueue()
self.send_all_hook = send_all_hook
self.wait_send_all_might_not_block_hook = wait_send_all_might_not_block_hook
self.close_hook = close_hook
async def send_all(self, data):
"""Places the given data into the object's internal buffer, and then
calls the :attr:`send_all_hook` (if any).
"""
# Execute two checkpoints so we have more of a chance to detect
# buggy user code that calls this twice at the same time.
with self._conflict_detector:
await _core.checkpoint()
await _core.checkpoint()
self._outgoing.put(data)
if self.send_all_hook is not None:
await self.send_all_hook()
async def wait_send_all_might_not_block(self):
"""Calls the :attr:`wait_send_all_might_not_block_hook` (if any), and
then returns immediately.
"""
# Execute two checkpoints so we have more of a chance to detect
# buggy user code that calls this twice at the same time.
with self._conflict_detector:
await _core.checkpoint()
await _core.checkpoint()
# check for being closed:
self._outgoing.put(b"")
if self.wait_send_all_might_not_block_hook is not None:
await self.wait_send_all_might_not_block_hook()
def close(self):
"""Marks this stream as closed, and then calls the :attr:`close_hook`
(if any).
"""
# XXX should this cancel any pending calls to the send_all_hook and
# wait_send_all_might_not_block_hook? Those are the only places where
# send_all and wait_send_all_might_not_block can be blocked.
#
# The way we set things up, send_all_hook is memory_stream_pump, and
# wait_send_all_might_not_block_hook is unset. memory_stream_pump is
# synchronous. So normally, send_all and wait_send_all_might_not_block
# cannot block at all.
self._outgoing.close()
if self.close_hook is not None:
self.close_hook()
async def aclose(self):
"""Same as :meth:`close`, but async."""
self.close()
await _core.checkpoint()
async def get_data(self, max_bytes=None):
"""Retrieves data from the internal buffer, blocking if necessary.
Args:
max_bytes (int or None): The maximum amount of data to
retrieve. None (the default) means to retrieve all the data
that's present (but still blocks until at least one byte is
available).
Returns:
If this stream has been closed, an empty bytearray. Otherwise, the
requested data.
"""
return await self._outgoing.get(max_bytes)
def get_data_nowait(self, max_bytes=None):
"""Retrieves data from the internal buffer, but doesn't block.
See :meth:`get_data` for details.
Raises:
trio.WouldBlock: if no data is available to retrieve.
"""
return self._outgoing.get_nowait(max_bytes)
class MemoryReceiveStream(ReceiveStream, metaclass=_util.Final):
"""An in-memory :class:`~trio.abc.ReceiveStream`.
Args:
receive_some_hook: An async function, or None. Called from
:meth:`receive_some`. Can do whatever you like.
close_hook: A synchronous function, or None. Called from :meth:`close`
and :meth:`aclose`. Can do whatever you like.
.. attribute:: receive_some_hook
close_hook
Both hooks are also exposed as attributes on the object, and you can
change them at any time.
"""
def __init__(self, receive_some_hook=None, close_hook=None):
self._conflict_detector = _util.ConflictDetector(
"another task is using this stream"
)
self._incoming = _UnboundedByteQueue()
self._closed = False
self.receive_some_hook = receive_some_hook
self.close_hook = close_hook
async def receive_some(self, max_bytes=None):
"""Calls the :attr:`receive_some_hook` (if any), and then retrieves
data from the internal buffer, blocking if necessary.
"""
# Execute two checkpoints so we have more of a chance to detect
# buggy user code that calls this twice at the same time.
with self._conflict_detector:
await _core.checkpoint()
await _core.checkpoint()
if self._closed:
raise _core.ClosedResourceError
if self.receive_some_hook is not None:
await self.receive_some_hook()
# self._incoming's closure state tracks whether we got an EOF.
# self._closed tracks whether we, ourselves, are closed.
# self.close() sends an EOF to wake us up and sets self._closed,
# so after we wake up we have to check self._closed again.
data = await self._incoming.get(max_bytes)
if self._closed:
raise _core.ClosedResourceError
return data
def close(self):
"""Discards any pending data from the internal buffer, and marks this
stream as closed.
"""
self._closed = True
self._incoming.close_and_wipe()
if self.close_hook is not None:
self.close_hook()
async def aclose(self):
"""Same as :meth:`close`, but async."""
self.close()
await _core.checkpoint()
def put_data(self, data):
"""Appends the given data to the internal buffer."""
self._incoming.put(data)
def put_eof(self):
"""Adds an end-of-file marker to the internal buffer."""
self._incoming.close()
def memory_stream_pump(memory_send_stream, memory_receive_stream, *, max_bytes=None):
"""Take data out of the given :class:`MemorySendStream`'s internal buffer,
and put it into the given :class:`MemoryReceiveStream`'s internal buffer.
Args:
memory_send_stream (MemorySendStream): The stream to get data from.
memory_receive_stream (MemoryReceiveStream): The stream to put data into.
max_bytes (int or None): The maximum amount of data to transfer in this
call, or None to transfer all available data.
Returns:
True if it successfully transferred some data, or False if there was no
data to transfer.
This is used to implement :func:`memory_stream_one_way_pair` and
:func:`memory_stream_pair`; see the latter's docstring for an example
of how you might use it yourself.
"""
try:
data = memory_send_stream.get_data_nowait(max_bytes)
except _core.WouldBlock:
return False
try:
if not data:
memory_receive_stream.put_eof()
else:
memory_receive_stream.put_data(data)
except _core.ClosedResourceError:
raise _core.BrokenResourceError("MemoryReceiveStream was closed")
return True
def memory_stream_one_way_pair():
"""Create a connected, pure-Python, unidirectional stream with infinite
buffering and flexible configuration options.
You can think of this as being a no-operating-system-involved
Trio-streamsified version of :func:`os.pipe` (except that :func:`os.pipe`
returns the streams in the wrong order – we follow the superior convention
that data flows from left to right).
Returns:
A tuple (:class:`MemorySendStream`, :class:`MemoryReceiveStream`), where
the :class:`MemorySendStream` has its hooks set up so that it calls
:func:`memory_stream_pump` from its
:attr:`~MemorySendStream.send_all_hook` and
:attr:`~MemorySendStream.close_hook`.
The end result is that data automatically flows from the
:class:`MemorySendStream` to the :class:`MemoryReceiveStream`. But you're
also free to rearrange things however you like. For example, you can
temporarily set the :attr:`~MemorySendStream.send_all_hook` to None if you
want to simulate a stall in data transmission. Or see
:func:`memory_stream_pair` for a more elaborate example.
"""
send_stream = MemorySendStream()
recv_stream = MemoryReceiveStream()
def pump_from_send_stream_to_recv_stream():
memory_stream_pump(send_stream, recv_stream)
async def async_pump_from_send_stream_to_recv_stream():
pump_from_send_stream_to_recv_stream()
send_stream.send_all_hook = async_pump_from_send_stream_to_recv_stream
send_stream.close_hook = pump_from_send_stream_to_recv_stream
return send_stream, recv_stream
def _make_stapled_pair(one_way_pair):
pipe1_send, pipe1_recv = one_way_pair()
pipe2_send, pipe2_recv = one_way_pair()
stream1 = StapledStream(pipe1_send, pipe2_recv)
stream2 = StapledStream(pipe2_send, pipe1_recv)
return stream1, stream2
def memory_stream_pair():
"""Create a connected, pure-Python, bidirectional stream with infinite
buffering and flexible configuration options.
This is a convenience function that creates two one-way streams using
:func:`memory_stream_one_way_pair`, and then uses
:class:`~trio.StapledStream` to combine them into a single bidirectional
stream.
This is like a no-operating-system-involved, Trio-streamsified version of
:func:`socket.socketpair`.
Returns:
A pair of :class:`~trio.StapledStream` objects that are connected so
that data automatically flows from one to the other in both directions.
After creating a stream pair, you can send data back and forth, which is
enough for simple tests::
left, right = memory_stream_pair()
await left.send_all(b"123")
assert await right.receive_some() == b"123"
await right.send_all(b"456")
assert await left.receive_some() == b"456"
But if you read the docs for :class:`~trio.StapledStream` and
:func:`memory_stream_one_way_pair`, you'll see that all the pieces
involved in wiring this up are public APIs, so you can adjust to suit the
requirements of your tests. For example, here's how to tweak a stream so
that data flowing from left to right trickles in one byte at a time (but
data flowing from right to left proceeds at full speed)::
left, right = memory_stream_pair()
async def trickle():
# left is a StapledStream, and left.send_stream is a MemorySendStream
# right is a StapledStream, and right.recv_stream is a MemoryReceiveStream
while memory_stream_pump(left.send_stream, right.recv_stream, max_bytes=1):
# Pause between each byte
await trio.sleep(1)
# Normally this send_all_hook calls memory_stream_pump directly without
# passing in a max_bytes. We replace it with our custom version:
left.send_stream.send_all_hook = trickle
And here's a simple test using our modified stream objects::
async def sender():
await left.send_all(b"12345")
await left.send_eof()
async def receiver():
async for data in right:
print(data)
async with trio.open_nursery() as nursery:
nursery.start_soon(sender)
nursery.start_soon(receiver)
By default, this will print ``b"12345"`` and then immediately exit; with
our trickle stream it instead sleeps 1 second, then prints ``b"1"``, then
sleeps 1 second, then prints ``b"2"``, etc.
Pro-tip: you can insert sleep calls (like in our example above) to
manipulate the flow of data across tasks... and then use
:class:`MockClock` and its :attr:`~MockClock.autojump_threshold`
functionality to keep your test suite running quickly.
If you want to stress test a protocol implementation, one nice trick is to
use the :mod:`random` module (preferably with a fixed seed) to move random
numbers of bytes at a time, and insert random sleeps in between them. You
can also set up a custom :attr:`~MemoryReceiveStream.receive_some_hook` if
you want to manipulate things on the receiving side, and not just the
sending side.
"""
return _make_stapled_pair(memory_stream_one_way_pair)
################################################################
# In-memory streams - Lockstep version
################################################################
class _LockstepByteQueue:
def __init__(self):
self._data = bytearray()
self._sender_closed = False
self._receiver_closed = False
self._receiver_waiting = False
self._waiters = _core.ParkingLot()
self._send_conflict_detector = _util.ConflictDetector(
"another task is already sending"
)
self._receive_conflict_detector = _util.ConflictDetector(
"another task is already receiving"
)
def _something_happened(self):
self._waiters.unpark_all()
# Always wakes up when one side is closed, because everyone always reacts
# to that.
async def _wait_for(self, fn):
while True:
if fn():
break
if self._sender_closed or self._receiver_closed:
break
await self._waiters.park()
await _core.checkpoint()
def close_sender(self):
self._sender_closed = True
self._something_happened()
def close_receiver(self):
self._receiver_closed = True
self._something_happened()
async def send_all(self, data):
with self._send_conflict_detector:
if self._sender_closed:
raise _core.ClosedResourceError
if self._receiver_closed:
raise _core.BrokenResourceError
assert not self._data
self._data += data
self._something_happened()
await self._wait_for(lambda: not self._data)
if self._sender_closed:
raise _core.ClosedResourceError
if self._data and self._receiver_closed:
raise _core.BrokenResourceError
async def wait_send_all_might_not_block(self):
with self._send_conflict_detector:
if self._sender_closed:
raise _core.ClosedResourceError
if self._receiver_closed:
await _core.checkpoint()
return
await self._wait_for(lambda: self._receiver_waiting)
if self._sender_closed:
raise _core.ClosedResourceError
async def receive_some(self, max_bytes=None):
with self._receive_conflict_detector:
# Argument validation
if max_bytes is not None:
max_bytes = operator.index(max_bytes)
if max_bytes < 1:
raise ValueError("max_bytes must be >= 1")
# State validation
if self._receiver_closed:
raise _core.ClosedResourceError
# Wake wait_send_all_might_not_block and wait for data
self._receiver_waiting = True
self._something_happened()
try:
await self._wait_for(lambda: self._data)
finally:
self._receiver_waiting = False
if self._receiver_closed:
raise _core.ClosedResourceError
# Get data, possibly waking send_all
if self._data:
# Neat trick: if max_bytes is None, then obj[:max_bytes] is
# the same as obj[:].
got = self._data[:max_bytes]
del self._data[:max_bytes]
self._something_happened()
return got
else:
assert self._sender_closed
return b""
class _LockstepSendStream(SendStream):
def __init__(self, lbq):
self._lbq = lbq
def close(self):
self._lbq.close_sender()
async def aclose(self):
self.close()
await _core.checkpoint()
async def send_all(self, data):
await self._lbq.send_all(data)
async def wait_send_all_might_not_block(self):
await self._lbq.wait_send_all_might_not_block()
class _LockstepReceiveStream(ReceiveStream):
def __init__(self, lbq):
self._lbq = lbq
def close(self):
self._lbq.close_receiver()
async def aclose(self):
self.close()
await _core.checkpoint()
async def receive_some(self, max_bytes=None):
return await self._lbq.receive_some(max_bytes)
def lockstep_stream_one_way_pair():
"""Create a connected, pure Python, unidirectional stream where data flows
in lockstep.
Returns:
A tuple
(:class:`~trio.abc.SendStream`, :class:`~trio.abc.ReceiveStream`).
This stream has *absolutely no* buffering. Each call to
:meth:`~trio.abc.SendStream.send_all` will block until all the given data
has been returned by a call to
:meth:`~trio.abc.ReceiveStream.receive_some`.
This can be useful for testing flow control mechanisms in an extreme case,
or for setting up "clogged" streams to use with
:func:`check_one_way_stream` and friends.
In addition to fulfilling the :class:`~trio.abc.SendStream` and
:class:`~trio.abc.ReceiveStream` interfaces, the return objects
also have a synchronous ``close`` method.
"""
lbq = _LockstepByteQueue()
return _LockstepSendStream(lbq), _LockstepReceiveStream(lbq)
def lockstep_stream_pair():
"""Create a connected, pure-Python, bidirectional stream where data flows
in lockstep.
Returns:
A tuple (:class:`~trio.StapledStream`, :class:`~trio.StapledStream`).
This is a convenience function that creates two one-way streams using
:func:`lockstep_stream_one_way_pair`, and then uses
:class:`~trio.StapledStream` to combine them into a single bidirectional
stream.
"""
return _make_stapled_pair(lockstep_stream_one_way_pair)