streams.py
python
__all__ = (
'StreamReader', 'StreamWriter', 'StreamReaderProtocol',
'open_connection', 'start_server')
import collections
import socket
import sys
import warnings
import weakref
if hasattr(socket, 'AF_UNIX'):
__all__ += ('open_unix_connection', 'start_unix_server')
from . import coroutines
from . import events
from . import exceptions
from . import format_helpers
from . import protocols
from .log import logger
from .tasks import sleep
_DEFAULT_LIMIT = 2 ** 16 # 64 KiB
async def open_connection(host=None, port=None, *,
limit=_DEFAULT_LIMIT, **kwds):
"""Обёртка для create_connection(), возвращающая пару (reader, writer).
Возвращаемый reader — экземпляр StreamReader; writer — экземпляр StreamWriter.
Аргументы — все обычные аргументы create_connection() за исключением protocol_factory; наиболее распространённые — позиционные host и port, с последующими различными необязательными именованными аргументами.
Дополнительные необязательные именованные аргументы: loop (для указания используемого экземпляра цикла событий) и limit (для установки лимита буфера, передаваемого в StreamReader).
(Если вы хотите настроить классы StreamReader и/или StreamReaderProtocol, просто скопируйте код — здесь нет ничего особенного, кроме удобства.)
"""
loop = events.get_running_loop()
reader = StreamReader(limit=limit, loop=loop)
protocol = StreamReaderProtocol(reader, loop=loop)
transport, _ = await loop.create_connection(
lambda: protocol, host, port, **kwds)
writer = StreamWriter(transport, protocol, reader, loop)
return reader, writer
async def start_server(client_connected_cb, host=None, port=None, *,
limit=_DEFAULT_LIMIT, **kwds):
"""Запустить сокетный сервер, вызывать обратный вызов для каждого подключившегося клиента.
Первый параметр, `client_connected_cb`, принимает два параметра:
client_reader, client_writer. client_reader — это объект StreamReader,
а client_writer — объект StreamWriter. Этот
параметр может быть как обычной функцией обратного вызова, так и корутиной;
если это корутина, она будет автоматически преобразована в
задачу.
Остальные аргументы — это все обычные аргументы для
loop.create_server() за исключением protocol_factory; наиболее распространены
позиционные host и port, за которыми следуют различные необязательные именованные аргументы.
Возвращаемое значение такое же, как у loop.create_server().
Дополнительный необязательный именованный аргумент — limit (для установки
лимита буфера, передаваемого в StreamReader).
Возвращаемое значение такое же, как у loop.create_server(), т.е.
объект Server, который можно использовать для остановки службы.
"""
loop = events.get_running_loop()
def factory():
reader = StreamReader(limit=limit, loop=loop)
protocol = StreamReaderProtocol(reader, client_connected_cb,
loop=loop)
return protocol
return await loop.create_server(factory, host, port, **kwds)
if hasattr(socket, 'AF_UNIX'):
# UNIX Domain Sockets поддерживаются на этой платформе
async def open_unix_connection(path=None, *,
limit=_DEFAULT_LIMIT, **kwds):
"""Аналогично `open_connection`, но работает с UNIX Domain Sockets."""
loop = events.get_running_loop()
reader = StreamReader(limit=limit, loop=loop)
protocol = StreamReaderProtocol(reader, loop=loop)
transport, _ = await loop.create_unix_connection(
lambda: protocol, path, **kwds)
writer = StreamWriter(transport, protocol, reader, loop)
return reader, writer
async def start_unix_server(client_connected_cb, path=None, *,
limit=_DEFAULT_LIMIT, **kwds):
"""Аналогично `start_server`, но работает с UNIX Domain Sockets."""
loop = events.get_running_loop()
def factory():
reader = StreamReader(limit=limit, loop=loop)
protocol = StreamReaderProtocol(reader, client_connected_cb,
loop=loop)
return protocol
return await loop.create_unix_server(factory, path, **kwds)
class FlowControlMixin(protocols.Protocol):
"""Многократно используемая логика управления потоком для StreamWriter.drain().
Реализует методы протокола pause_writing(), resume_writing() и connection_lost(). Если подкласс переопределяет их, он должен вызывать методы родительского класса.
StreamWriter.drain() должен ожидать корутину _drain_helper().
"""
def __init__(self, loop=None):
if loop is None:
self._loop = events.get_event_loop()
else:
self._loop = loop
self._paused = False
self._drain_waiters = collections.deque()
self._connection_lost = False
def pause_writing(self):
assert not self._paused
self._paused = True
if self._loop.get_debug():
logger.debug("%r pauses writing", self)
def resume_writing(self):
assert self._paused
self._paused = False
if self._loop.get_debug():
logger.debug("%r resumes writing", self)
for waiter in self._drain_waiters:
if not waiter.done():
waiter.set_result(None)
def connection_lost(self, exc):
self._connection_lost = True
# Разбудить writer(s), если он(и) приостановлен(ы).
if not self._paused:
return
for waiter in self._drain_waiters:
if not waiter.done():
if exc is None:
waiter.set_result(None)
else:
waiter.set_exception(exc)
async def _drain_helper(self):
if self._connection_lost:
raise ConnectionResetError('Connection lost')
if not self._paused:
return
waiter = self._loop.create_future()
self._drain_waiters.append(waiter)
try:
await waiter
finally:
self._drain_waiters.remove(waiter)
def _get_close_waiter(self, stream):
raise NotImplementedError
class StreamReaderProtocol(FlowControlMixin, protocols.Protocol):
"""Вспомогательный класс для адаптации между Protocol и StreamReader.
(Это вспомогательный класс, а не подкласс StreamReader от Protocol,
потому что StreamReader имеет другие потенциальные применения,
и чтобы предотвратить случайный вызов пользователем StreamReader
неподходящих методов протокола.)
"""
_source_traceback = None
def __init__(self, stream_reader, client_connected_cb=None, loop=None):
super().__init__(loop=loop)
if stream_reader is not None:
self._stream_reader_wr = weakref.ref(stream_reader)
self._source_traceback = stream_reader._source_traceback
else:
self._stream_reader_wr = None
if client_connected_cb is not None:
# Это поток данных, созданный функцией `create_server()`.
# Сохраняйте строгую ссылку на reader до тех пор, пока соединение
# не будет установлено.
self._strong_reader = stream_reader
self._reject_connection = False
self._task = None
self._transport = None
self._client_connected_cb = client_connected_cb
self._over_ssl = False
self._closed = self._loop.create_future()
@property
def _stream_reader(self):
if self._stream_reader_wr is None:
return None
return self._stream_reader_wr()
def _replace_transport(self, transport):
self._transport = transport
self._over_ssl = transport.get_extra_info('sslcontext') is not None
def connection_made(self, transport):
if self._reject_connection:
context = {
'message': ('An open stream was garbage collected prior to '
'establishing network connection; '
'call "stream.close()" explicitly.')
}
if self._source_traceback:
context['source_traceback'] = self._source_traceback
self._loop.call_exception_handler(context)
transport.abort()
return
self._transport = transport
reader = self._stream_reader
if reader is not None:
reader.set_transport(transport)
self._over_ssl = transport.get_extra_info('sslcontext') is not None
if self._client_connected_cb is not None:
writer = StreamWriter(transport, self, reader, self._loop)
res = self._client_connected_cb(reader, writer)
if coroutines.iscoroutine(res):
def callback(task):
if task.cancelled():
transport.close()
return
exc = task.exception()
if exc is not None:
self._loop.call_exception_handler({
'message': 'Unhandled exception in client_connected_cb',
'exception': exc,
'transport': transport,
})
transport.close()
self._task = self._loop.create_task(res)
self._task.add_done_callback(callback)
self._strong_reader = None
def connection_lost(self, exc):
reader = self._stream_reader
if reader is not None:
if exc is None:
reader.feed_eof()
else:
reader.set_exception(exc)
if not self._closed.done():
if exc is None:
self._closed.set_result(None)
else:
self._closed.set_exception(exc)
super().connection_lost(exc)
self._stream_reader_wr = None
self._task = None
self._transport = None
def data_received(self, data):
reader = self._stream_reader
if reader is not None:
reader.feed_data(data)
def eof_received(self):
reader = self._stream_reader
if reader is not None:
reader.feed_eof()
if self._over_ssl:
# Предотвратить предупреждение в SSLProtocol.eof_received:
# "возврат true из eof_received()
# не имеет эффекта при использовании ssl"
return False
return True
def _get_close_waiter(self, stream):
return self._closed
def __del__(self):
# Предотвратить сообщения о необработанных исключениях.
# Лучше, чем хак self._closed._log_traceback = False
try:
closed = self._closed
except AttributeError:
pass # ошибка конструктора
else:
if closed.done() and not closed.cancelled():
closed.exception()
class StreamWriter:
"""Оборачивает Transport.
Предоставляет write(), writelines(), [can_]write_eof(),
get_extra_info() и close(). Добавляет drain(), который возвращает
опциональный Future, на котором можно ожидать управления потоком.
Также добавляет свойство transport, которое напрямую ссылается на Transport.
"""
def __init__(self, transport, protocol, reader, loop):
self._transport = transport
self._protocol = protocol
# drain() ожидает, что reader имеет метод exception()
assert reader is None or isinstance(reader, StreamReader)
self._reader = reader
self._loop = loop
self._complete_fut = self._loop.create_future()
self._complete_fut.set_result(None)
def __repr__(self):
info = [self.__class__.__name__, f'transport={self._transport!r}']
if self._reader is not None:
info.append(f'reader={self._reader!r}')
return '<{}>'.format(' '.join(info))
@property
def transport(self):
return self._transport
def write(self, data):
self._transport.write(data)
def writelines(self, data):
self._transport.writelines(data)
def write_eof(self):
return self._transport.write_eof()
def can_write_eof(self):
return self._transport.can_write_eof()
def close(self):
return self._transport.close()
def is_closing(self):
return self._transport.is_closing()
async def wait_closed(self):
await self._protocol._get_close_waiter(self)
def get_extra_info(self, name, default=None):
return self._transport.get_extra_info(name, default)
async def drain(self):
"""Сбросить буфер записи.
Предполагаемое использование:
w.write(data)
await w.drain()
"""
if self._reader is not None:
exc = self._reader.exception()
if exc is not None:
raise exc
if self._transport.is_closing():
# Ожидать вызова protocol.connection_lost()
# Вызвать ошибку закрытия соединения, если есть,
# иначе ConnectionResetError
# Уступите управление циклу событий, чтобы connection_lost() могла быть
# вызвана. Без этого _drain_helper() вернула бы
# немедленно, и код, который вызывает
# write(...); await drain()
# в цикле никогда бы не вызвал connection_lost(), поэтому он
# не увидел бы ошибку при закрытии сокета.
await sleep(0)
await self._protocol._drain_helper()
async def start_tls(self, sslcontext, *,
server_hostname=None,
ssl_handshake_timeout=None,
ssl_shutdown_timeout=None):
"""Обновить существующее потоковое соединение до TLS."""
server_side = self._protocol._client_connected_cb is not None
protocol = self._protocol
await self.drain()
new_transport = await self._loop.start_tls( # type: ignore
self._transport, protocol, sslcontext,
server_side=server_side, server_hostname=server_hostname,
ssl_handshake_timeout=ssl_handshake_timeout,
ssl_shutdown_timeout=ssl_shutdown_timeout)
self._transport = new_transport
protocol._replace_transport(new_transport)
def __del__(self, warnings=warnings):
if not self._transport.is_closing():
if self._loop.is_closed():
warnings.warn("loop is closed", ResourceWarning)
else:
self.close()
warnings.warn(f"unclosed {self!r}", ResourceWarning)
class StreamReader:
_source_traceback = None
def __init__(self, limit=_DEFAULT_LIMIT, loop=None):
# Ограничение длины строки является функцией безопасности;
# оно также служит половиной лимита буфера.
if limit <= 0:
raise ValueError('Limit cannot be <= 0')
self._limit = limit
if loop is None:
self._loop = events.get_event_loop()
else:
self._loop = loop
self._buffer = bytearray()
self._eof = False # Указывает, завершена ли работа.
self._waiter = None # Future, используемый _wait_for_data()
self._exception = None
self._transport = None
self._paused = False
if self._loop.get_debug():
self._source_traceback = format_helpers.extract_stack(
sys._getframe(1))
def __repr__(self):
info = ['StreamReader']
if self._buffer:
info.append(f'{len(self._buffer)} bytes')
if self._eof:
info.append('eof')
if self._limit != _DEFAULT_LIMIT:
info.append(f'limit={self._limit}')
if self._waiter:
info.append(f'waiter={self._waiter!r}')
if self._exception:
info.append(f'exception={self._exception!r}')
if self._transport:
info.append(f'transport={self._transport!r}')
if self._paused:
info.append('paused')
return '<{}>'.format(' '.join(info))
def exception(self):
return self._exception
def set_exception(self, exc):
self._exception = exc
waiter = self._waiter
if waiter is not None:
self._waiter = None
if not waiter.cancelled():
waiter.set_exception(exc)
def _wakeup_waiter(self):
"""Пробуждает функции read*(), ожидающие данные или EOF."""
waiter = self._waiter
if waiter is not None:
self._waiter = None
if not waiter.cancelled():
waiter.set_result(None)
def set_transport(self, transport):
assert self._transport is None, 'Transport already set'
self._transport = transport
def _maybe_resume_transport(self):
if self._paused and len(self._buffer) <= self._limit:
self._paused = False
self._transport.resume_reading()
def feed_eof(self):
self._eof = True
self._wakeup_waiter()
def at_eof(self):
"""Возвращает True, если буфер пуст и был вызван 'feed_eof'."""
return self._eof and not self._buffer
def feed_data(self, data):
assert not self._eof, 'feed_data after feed_eof'
if not data:
return
self._buffer.extend(data)
self._wakeup_waiter()
if (self._transport is not None and
not self._paused and
len(self._buffer) > 2 * self._limit):
try:
self._transport.pause_reading()
except NotImplementedError:
# Транспорт не может быть приостановлен.
# Придется просто буферизировать все данные.
# Забываем транспорт, чтобы не продолжать попытки.
self._transport = None
else:
self._paused = True
async def _wait_for_data(self, func_name):
"""Ожидать, пока не будет вызван feed_data() или feed_eof().
Если поток был приостановлен, автоматически возобновить его.
"""
# StreamReader использует future для связывания метода feed_data() протокола
# с корутиной чтения. Одновременное выполнение двух корутин чтения
# приведет к неожиданному поведению. Невозможно будет узнать
# какая корутина получит следующие данные.
if self._waiter is not None:
raise RuntimeError(
f'{func_name}() called while another coroutine is '
f'already waiting for incoming data')
assert not self._eof, '_wait_for_data after EOF'
# Ожидание данных в приостановленном состоянии приведет к взаимоблокировке, поэтому предотвратите это.
# Это важно для readexactly(n) в случае, когда n > self._limit.
if self._paused:
self._paused = False
self._transport.resume_reading()
self._waiter = self._loop.create_future()
try:
await self._waiter
finally:
self._waiter = None
async def readline(self):
r"""Прочитать часть данных из потока до тех пор, пока не будет найден символ новой строки (b'\n').
В случае успеха возвращает часть, оканчивающуюся символом новой строки. Если из-за EOF
можно прочитать только неполную строку, возвращается неполная строка без завершающего
символа новой строки. Когда EOF достигнут, а байты не прочитаны,
возвращается пустой объект bytes.
Если превышен лимит, будет вызвано ValueError. В этом случае, если
символ новой строки найден, полная строка, включая символ новой строки, будет удалена
из внутреннего буфера. В противном случае внутренний буфер будет очищен.
Лимит сравнивается с частью строки без символа новой строки.
Если поток был приостановлен, эта функция автоматически возобновит его, если
необходимо.
"""
sep = b'\n'
seplen = len(sep)
try:
line = await self.readuntil(sep)
except exceptions.IncompleteReadError as e:
return e.partial
except exceptions.LimitOverrunError as e:
if self._buffer.startswith(sep, e.consumed):
del self._buffer[:e.consumed + seplen]
else:
self._buffer.clear()
self._maybe_resume_transport()
raise ValueError(e.args[0])
return line
async def readuntil(self, separator=b'\n'):
"""Прочитать данные из потока до тех пор, пока не будет найден ``separator``.
В случае успеха данные и разделитель будут удалены из
внутреннего буфера (потреблены). Возвращаемые данные будут включать
разделитель в конце.
Для проверки результата используется настроенный лимит потока. Лимит устанавливает
максимальную длину возвращаемых данных, не считая разделителя.
Если произошел EOF и полный разделитель все еще не найден,
будет вызвано исключение IncompleteReadError, и внутренний
буфер будет сброшен. Атрибут IncompleteReadError.partial
может содержать частичный разделитель.
Если данные не могут быть прочитаны из-за превышения лимита,
будет вызвано исключение LimitOverrunError, и данные
останутся во внутреннем буфере, чтобы их можно было прочитать снова.
``separator`` также может быть кортежем разделителей. В этом
случае возвращаемое значение будет кратчайшим возможным, имеющим любой
разделитель в качестве суффикса. Для целей LimitOverrunError
кратчайший возможный разделитель считается тем, который совпал.
"""
if isinstance(separator, tuple):
# Гарантирует, что выигрывает кратчайшее совпадение
separator = sorted(separator, key=len)
else:
separator = [separator]
if not separator:
raise ValueError('Separator should contain at least one element')
min_seplen = len(separator[0])
max_seplen = len(separator[-1])
if min_seplen == 0:
raise ValueError('Separator should be at least one-byte string')
if self._exception is not None:
raise self._exception
# Потребляет весь буфер, кроме последних байтов, длина которых
# на единицу меньше max_seplen. Давайте проверим крайние случаи с
# separator[-1]='SEPARATOR':
# * мы получили почти полный разделитель (без последнего
# байта). т.е. buffer='some textSEPARATO'. В этом случае мы
# можно безопасно обработать max_seplen - 1 байт.
# * последний байт буфера является первым байтом разделителя, т.е.
# buffer='abcdefghijklmnopqrS'. Мы можем безопасно потребить
# всё, кроме последнего байта, но для этого требуется
# анализировать байты буфера, совпадающие с частичным разделителем.
# Это медленно и/или требует конечного автомата. В этом случае наша
# реализация не оптимальна, так как требует повторного сканирования
# данных, которые, как известно, не относятся к разделителю. В
# реальном мире разделитель не будет настолько длинным, чтобы заметить
# проблемы производительности. Даже при чтении MIME-кодированных
# сообщений :)
# `offset` — это количество байт от начала буфера,
# где нет ни одного вхождения `separator`.
offset = 0
# Цикл до тех пор, пока не найдем `separator` в буфере, не превысим размер буфера,
# или не произойдёт EOF.
while True:
buflen = len(self._buffer)
# Проверяем, достаточно ли теперь данных в буфере для кратчайшего
# разделителя, чтобы поместиться.
if buflen - offset >= min_seplen:
match_start = None
match_end = None
for sep in separator:
isep = self._buffer.find(sep, offset)
if isep != -1:
# `separator` находится в буфере. `match_start` и
# `match_end` будут использованы позже для извлечения
# данных.
end = isep + len(sep)
if match_end is None or end < match_end:
match_end = end
match_start = isep
if match_end is not None:
break
# смотрите комментарий выше для объяснения.
offset = max(0, buflen + 1 - max_seplen)
if offset > self._limit:
raise exceptions.LimitOverrunError(
'Separator is not found, and chunk exceed the limit',
offset)
# Полное сообщение (с полным разделителем) может присутствовать в буфере
# даже когда установлен флаг EOF. Это может произойти, когда последний фрагмент
# добавляет данные, из-за которых находится разделитель. Поэтому мы проверяем
# EOF *после* проверки буфера.
if self._eof:
chunk = self._buffer.take_bytes()
raise exceptions.IncompleteReadError(chunk, None)
# _wait_for_data() возобновит чтение, если поток был приостановлен.
await self._wait_for_data('readuntil')
if match_start > self._limit:
raise exceptions.LimitOverrunError(
'Separator is found, but chunk is longer than limit', match_start)
chunk = self._buffer.take_bytes(match_end)
self._maybe_resume_transport()
return chunk
async def read(self, n=-1):
"""Читает до `n` байт из потока.
Если `n` не указано или равно -1,
читает до EOF, затем возвращает все прочитанные байты.
Если был получен EOF и внутренний буфер пуст,
возвращает пустой объект bytes.
Если `n` равно 0, немедленно возвращает пустой объект bytes.
Если `n` положительное, возвращает не более `n` доступных байт
как только хотя бы 1 байт становится доступен во внутреннем буфере.
Если EOF получен до того, как прочитан хотя бы один байт, возвращает пустой
объект bytes.
Возвращаемое значение не ограничено лимитом, установленным при создании
потока.
Если поток был приостановлен, эта функция автоматически возобновит его при
необходимости.
"""
if self._exception is not None:
raise self._exception
if n == 0:
return b''
if n < 0:
# Раньше это просто зацикливалось, создавая новый waiter в надежде
# собрать все в self._buffer, но это бы
# привело к взаимоблокировке, если подпроцесс отправляет больше, чем self.limit
# байт. Поэтому просто вызываем self.read(self._limit) до EOF.
joined = bytearray()
while block := await self.read(self._limit):
joined += block
return joined.take_bytes()
if not self._buffer and not self._eof:
await self._wait_for_data('read')
# Это будет работать правильно, даже если буфер содержит меньше n байт
data = self._buffer.take_bytes(min(len(self._buffer), n))
self._maybe_resume_transport()
return data
async def readexactly(self, n):
"""Читает ровно `n` байт.
Возбуждает IncompleteReadError, если EOF достигнут до того, как может быть прочитано `n` байт.
Атрибут IncompleteReadError.partial исключения будет содержать частично прочитанные байты.
если n равно нулю, возвращает пустой объект bytes.
Возвращаемое значение не ограничено лимитом, установленным при создании
потока.
Если поток был приостановлен, эта функция автоматически возобновит его при
необходимости.
"""
if n < 0:
raise ValueError('readexactly size can not be less than zero')
if self._exception is not None:
raise self._exception
if n == 0:
return b''
while len(self._buffer) < n:
if self._eof:
incomplete = self._buffer.take_bytes()
raise exceptions.IncompleteReadError(incomplete, n)
await self._wait_for_data('readexactly')
data = self._buffer.take_bytes(n)
self._maybe_resume_transport()
return data
def __aiter__(self):
return self
async def __anext__(self):
val = await self.readline()
if val == b'':
raise StopAsyncIteration
return val