async def send_all(sock, data):
total_sent = 0
while total_sent < len(data):
sent = await sock.send(data[total_sent:])
if sent == 0:
raise trio.BrokenResourceError('socket connection broken')
total_sent += sent
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()
def _translate_socket_errors_to_stream_errors():
try:
yield
except OSError as exc:
if exc.errno in _closed_stream_errnos:
raise trio.ClosedResourceError(
"this socket was already closed") from None
else:
raise trio.BrokenResourceError(
"socket connection broken: {}".format(exc)) from exc
def _translate_socket_errors_to_stream_errors():
try:
yield
except OSError as exc:
if exc.errno in {errno.EBADF, errno.ENOTSOCK}:
# EBADF on Unix, ENOTSOCK on Windows
raise trio.ClosedResourceError(
"this socket was already closed") from None
else:
raise trio.BrokenResourceError(
"socket connection broken: {}".format(exc)) from exc
async def aclose(self):
if self._closed:
await trio.hazmat.checkpoint()
return
self._closed = True
for task in self._tasks:
trio.hazmat.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.hazmat.reschedule(task, Error(trio.BrokenResourceError()))
self._state.send_tasks.clear()
self._state.data.clear()
await trio.hazmat.checkpoint()
async def listen(self, table: str, *symbols: Optional[Sequence[str]]):
"""
Subscribe to a channel and optionally one or more specific symbols.
Returns an async generator that yields messages from the subscribed channel.
"""
if self._websocket.closed is not None:
raise trio.BrokenResourceError('Connection is closed.')
listeners = [(table, )] if not symbols else [(table, symbol)
for symbol in symbols]
args = []
for listener in listeners:
if self._subscriptions[listener] == 0:
args.append(listener[0] if not symbols else ':'.join(listener))
self._subscriptions[listener] += 1
await self._send_channel.send({'op': 'subscribe', 'args': args})
async with self._pipeline.tap() as aiter:
async for item, item_symbol, item_table, _ in aiter:
# Lock list of listeners while sending
if item_table == table and (not symbols
or item_symbol in symbols):
yield item
log.debug('Listener detached from table: %s, symbol: %s', table,
symbols)
if self._websocket.closed:
return
args = []
for listener in listeners:
self._subscriptions[listener] -= 1
if self._subscriptions[listener] == 0:
log.debug(
'No more listeners on table: %s, symbol: %s. Unsubscribing.',
table, symbols)
args.append(listener[0] if not symbols else ':'.join(listener))
if args:
await self._send_channel.send({'op': 'unsubscribe', 'args': args})
async def _broken_stream(*args, **kwargs):
raise trio.BrokenResourceError()
