streams.py
python
"""Вопросы, связанные с потоками данных."""
__all__ = ['StreamReader', 'StreamWriter', 'StreamReaderProtocol',
'open_connection', 'start_server',
'IncompleteReadError',
'LimitOverrunError',
]
import socket
if hasattr(socket, 'AF_UNIX'):
__all__.extend(['open_unix_connection', 'start_unix_server'])
from . import coroutines
from . import compat
from . import events
from . import protocols
from .coroutines import coroutine
from .log import logger
_DEFAULT_LIMIT = 2 ** 16
class IncompleteReadError(EOFError):
"""
Ошибка неполного чтения. Атрибуты:
- partial: строка байтов, прочитанная до достижения конца потока
- expected: общее количество ожидаемых байтов (или None, если неизвестно)
"""
def __init__(self, partial, expected):
super().__init__("%d bytes read on a total of %r expected bytes"
% (len(partial), expected))
self.partial = partial
self.expected = expected
def __reduce__(self):
return type(self), (self.partial, self.expected)
class LimitOverrunError(Exception):
"""Достигнут лимит буфера во время поиска разделителя.
Атрибуты:
- consumed: общее количество байтов, которые необходимо потребить.
"""
def __init__(self, message, consumed):
super().__init__(message)
self.consumed = consumed
def __reduce__(self):
return type(self), (self.args[0], self.consumed)
@coroutine
def open_connection(host=None, port=None, *,
loop=None, limit=_DEFAULT_LIMIT, **kwds):
"""Обёртка для create_connection(), возвращающая пару (reader, writer).
Возвращаемый reader — экземпляр StreamReader; writer — экземпляр StreamWriter.
Аргументы — все обычные аргументы create_connection() за исключением protocol_factory; наиболее распространённые — позиционные host и port, с последующими различными необязательными именованными аргументами.
Дополнительные необязательные именованные аргументы: loop (для указания используемого экземпляра цикла событий) и limit (для установки лимита буфера, передаваемого в StreamReader).
(Если вы хотите настроить классы StreamReader и/или StreamReaderProtocol, просто скопируйте код — здесь нет ничего особенного, кроме удобства.)
"""
if loop is None:
loop = events.get_event_loop()
reader = StreamReader(limit=limit, loop=loop)
protocol = StreamReaderProtocol(reader, loop=loop)
transport, _ = yield from loop.create_connection(
lambda: protocol, host, port, **kwds)
writer = StreamWriter(transport, protocol, reader, loop)
return reader, writer
@coroutine
def start_server(client_connected_cb, host=None, port=None, *,
loop=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().
Дополнительные необязательные именованные аргументы: loop (для указания используемого экземпляра цикла событий) и limit (для установки лимита буфера, передаваемого в StreamReader).
Возвращаемое значение то же, что и у loop.create_server(), т.е. объект Server, который можно использовать для остановки службы.
"""
if loop is None:
loop = events.get_event_loop()
def factory():
reader = StreamReader(limit=limit, loop=loop)
protocol = StreamReaderProtocol(reader, client_connected_cb,
loop=loop)
return protocol
return (yield from loop.create_server(factory, host, port, **kwds))
if hasattr(socket, 'AF_UNIX'):
# UNIX Domain Sockets поддерживаются на этой платформе
@coroutine
def open_unix_connection(path=None, *,
loop=None, limit=_DEFAULT_LIMIT, **kwds):
"""Аналогично `open_connection`, но работает с UNIX Domain Sockets."""
if loop is None:
loop = events.get_event_loop()
reader = StreamReader(limit=limit, loop=loop)
protocol = StreamReaderProtocol(reader, loop=loop)
transport, _ = yield from loop.create_unix_connection(
lambda: protocol, path, **kwds)
writer = StreamWriter(transport, protocol, reader, loop)
return reader, writer
@coroutine
def start_unix_server(client_connected_cb, path=None, *,
loop=None, limit=_DEFAULT_LIMIT, **kwds):
"""Аналогично `start_server`, но работает с UNIX Domain Sockets."""
if loop is None:
loop = events.get_event_loop()
def factory():
reader = StreamReader(limit=limit, loop=loop)
protocol = StreamReaderProtocol(reader, client_connected_cb,
loop=loop)
return protocol
return (yield from loop.create_unix_server(factory, path, **kwds))
class FlowControlMixin(protocols.Protocol):
"""Многоразовая логика управления потоком для StreamWriter.drain().
Реализует методы протокола pause_writing(), resume_reading() и 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_waiter = None
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)
waiter = self._drain_waiter
if waiter is not None:
self._drain_waiter = None
if not waiter.done():
waiter.set_result(None)
def connection_lost(self, exc):
self._connection_lost = True
# Разбудить writer, если он приостановлен.
if not self._paused:
return
waiter = self._drain_waiter
if waiter is None:
return
self._drain_waiter = None
if waiter.done():
return
if exc is None:
waiter.set_result(None)
else:
waiter.set_exception(exc)
@coroutine
def _drain_helper(self):
if self._connection_lost:
raise ConnectionResetError('Connection lost')
if not self._paused:
return
waiter = self._drain_waiter
assert waiter is None or waiter.cancelled()
waiter = self._loop.create_future()
self._drain_waiter = waiter
yield from waiter
class StreamReaderProtocol(FlowControlMixin, protocols.Protocol):
"""Вспомогательный класс для адаптации между Protocol и StreamReader.
(Это вспомогательный класс, а не подкласс StreamReader от Protocol,
потому что StreamReader имеет другие потенциальные применения,
и чтобы предотвратить случайный вызов пользователем StreamReader
неподходящих методов протокола.)
"""
def __init__(self, stream_reader, client_connected_cb=None, loop=None):
super().__init__(loop=loop)
self._stream_reader = stream_reader
self._stream_writer = None
self._client_connected_cb = client_connected_cb
self._over_ssl = False
def connection_made(self, transport):
self._stream_reader.set_transport(transport)
self._over_ssl = transport.get_extra_info('sslcontext') is not None
if self._client_connected_cb is not None:
self._stream_writer = StreamWriter(transport, self,
self._stream_reader,
self._loop)
res = self._client_connected_cb(self._stream_reader,
self._stream_writer)
if coroutines.iscoroutine(res):
self._loop.create_task(res)
def connection_lost(self, exc):
if self._stream_reader is not None:
if exc is None:
self._stream_reader.feed_eof()
else:
self._stream_reader.set_exception(exc)
super().connection_lost(exc)
self._stream_reader = None
self._stream_writer = None
def data_received(self, data):
self._stream_reader.feed_data(data)
def eof_received(self):
self._stream_reader.feed_eof()
if self._over_ssl:
# Предотвратить предупреждение в SSLProtocol.eof_received:
# "возврат true из eof_received()
# не имеет эффекта при использовании ssl"
return False
return True
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
def __repr__(self):
info = [self.__class__.__name__, 'transport=%r' % self._transport]
if self._reader is not None:
info.append('reader=%r' % self._reader)
return '<%s>' % ' '.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 get_extra_info(self, name, default=None):
return self._transport.get_extra_info(name, default)
@coroutine
def drain(self):
"""Сбросить буфер записи.
Предполагаемое использование — написать
w.write(data)
yield from w.drain()
"""
if self._reader is not None:
exc = self._reader.exception()
if exc is not None:
raise exc
if self._transport is not None:
if self._transport.is_closing():
# Уступите управление циклу событий, чтобы connection_lost() могла быть
# вызвана. Без этого _drain_helper() вернула бы
# немедленно, и код, который вызывает
# write(...); yield from drain()
# в цикле никогда бы не вызвал connection_lost(), поэтому он
# не увидел бы ошибку при закрытии сокета.
yield
yield from self._protocol._drain_helper()
class StreamReader:
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
def __repr__(self):
info = ['StreamReader']
if self._buffer:
info.append('%d bytes' % len(self._buffer))
if self._eof:
info.append('eof')
if self._limit != _DEFAULT_LIMIT:
info.append('l=%d' % self._limit)
if self._waiter:
info.append('w=%r' % self._waiter)
if self._exception:
info.append('e=%r' % self._exception)
if self._transport:
info.append('t=%r' % self._transport)
if self._paused:
info.append('paused')
return '<%s>' % ' '.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
@coroutine
def _wait_for_data(self, func_name):
"""Ожидать, пока не будет вызван feed_data() или feed_eof().
Если поток был приостановлен, автоматически возобновить его.
"""
# StreamReader использует future для связывания метода feed_data() протокола
# с корутиной чтения. Одновременное выполнение двух корутин чтения
# приведет к неожиданному поведению. Невозможно будет узнать
# какая корутина получит следующие данные.
if self._waiter is not None:
raise RuntimeError('%s() called while another coroutine is '
'already waiting for incoming data' % func_name)
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:
yield from self._waiter
finally:
self._waiter = None
@coroutine
def readline(self):
"""Читает порцию данных из потока до нахождения символа новой строки (b'\n').
В случае успеха возвращает порцию, заканчивающуюся символом новой строки. Если из-за EOF можно прочитать только неполную строку, возвращает неполную строку без завершающего символа новой строки. Когда EOF достигнут, а байты не были прочитаны, возвращается пустой объект bytes.
Если достигнут лимит, будет возбуждено исключение ValueError. В этом случае, если символ новой строки найден, полная строка вместе с символом новой строки будет удалена из внутреннего буфера. В противном случае внутренний буфер будет очищен. Лимит сравнивается с частью строки без символа новой строки.
Если поток был приостановлен, эта функция автоматически возобновит его при необходимости.
"""
sep = b'\n'
seplen = len(sep)
try:
line = yield from self.readuntil(sep)
except IncompleteReadError as e:
return e.partial
except 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
@coroutine
def readuntil(self, separator=b'\n'):
"""Читает данные из потока до нахождения ``separator``.
В случае успеха данные и разделитель будут удалены из внутреннего буфера (потреблены). Возвращаемые данные будут включать разделитель в конце.
Настроенный лимит потока используется для проверки результата. Лимит задает максимальную длину данных, которые могут быть возвращены, не считая разделителя.
Если произошел EOF, а полный разделитель все еще не найден, будет возбуждено исключение IncompleteReadError, и внутренний буфер будет сброшен. Атрибут IncompleteReadError.partial может содержать часть разделителя.
Если данные не могут быть прочитаны из-за превышения лимита, будет возбуждено исключение LimitOverrunError, и данные останутся во внутреннем буфере, чтобы их можно было прочитать снова.
"""
seplen = len(separator)
if seplen == 0:
raise ValueError('Separator should be at least one-byte string')
if self._exception is not None:
raise self._exception
# Потребляет весь буфер, кроме последних байтов, длина которых
# на единицу меньше seplen. Рассмотрим граничные случаи с
# separator='SEPARATOR':
# * мы получили почти полный разделитель (без последнего
# байта). т.е. buffer='some textSEPARATO'. В этом случае мы
# можем безопасно потребить len(separator) - 1 байт.
# * последний байт буфера является первым байтом разделителя, т.е.
# buffer='abcdefghijklmnopqrS'. Мы можем безопасно потребить
# всё, кроме последнего байта, но для этого требуется
# анализировать байты буфера, совпадающие с частичным разделителем.
# Это медленно и/или требует конечного автомата. В этом случае наша
# реализация не оптимальна, так как требует повторного сканирования
# данных, которые, как известно, не относятся к разделителю. В
# реальном мире разделитель не будет настолько длинным, чтобы заметить
# проблемы производительности. Даже при чтении MIME-кодированных
# сообщений :)
# `offset` — это количество байт от начала буфера,
# где нет вхождения `separator`.
offset = 0
# Цикл до тех пор, пока мы не найдём `separator` в буфере, не превысим размер буфера,
# или не произойдёт EOF.
while True:
buflen = len(self._buffer)
# Проверить, достаточно ли сейчас данных в буфере, чтобы `separator`
# поместился.
if buflen - offset >= seplen:
isep = self._buffer.find(separator, offset)
if isep != -1:
# `separator` находится в буфере. `isep` будет использован позже
# для извлечения данных.
break
# смотрите комментарий выше для объяснения.
offset = buflen + 1 - seplen
if offset > self._limit:
raise LimitOverrunError(
'Separator is not found, and chunk exceed the limit',
offset)
# Полное сообщение (с полным разделителем) может присутствовать в буфере
# даже когда установлен флаг EOF. Это может произойти, когда последний фрагмент
# добавляет данные, из-за которых находится разделитель. Поэтому мы проверяем
# EOF *после* проверки буфера.
if self._eof:
chunk = bytes(self._buffer)
self._buffer.clear()
raise IncompleteReadError(chunk, None)
# _wait_for_data() возобновит чтение, если поток был приостановлен.
yield from self._wait_for_data('readuntil')
if isep > self._limit:
raise LimitOverrunError(
'Separator is found, but chunk is longer than limit', isep)
chunk = self._buffer[:isep + seplen]
del self._buffer[:isep + seplen]
self._maybe_resume_transport()
return bytes(chunk)
@coroutine
def read(self, n=-1):
"""Прочитать до `n` байт из потока данных.
Если n не указано или равно -1, читать до EOF и вернуть все прочитанные
байты. Если EOF был получен и внутренний буфер пуст, вернуть
пустой объект bytes.
Если n равно нулю, немедленно вернуть пустой объект bytes.
Если n положительно, эта функция пытается прочитать `n` байт и может вернуть
меньше или равно запрошенному, но как минимум один байт. Если 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.
blocks = []
while True:
block = yield from self.read(self._limit)
if not block:
break
blocks.append(block)
return b''.join(blocks)
if not self._buffer and not self._eof:
yield from self._wait_for_data('read')
# Это будет работать правильно, даже если буфер содержит меньше n байт
data = bytes(self._buffer[:n])
del self._buffer[:n]
self._maybe_resume_transport()
return data
@coroutine
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 = bytes(self._buffer)
self._buffer.clear()
raise IncompleteReadError(incomplete, n)
yield from self._wait_for_data('readexactly')
if len(self._buffer) == n:
data = bytes(self._buffer)
self._buffer.clear()
else:
data = bytes(self._buffer[:n])
del self._buffer[:n]
self._maybe_resume_transport()
return data
if compat.PY35:
@coroutine
def __aiter__(self):
return self
@coroutine
def __anext__(self):
val = yield from self.readline()
if val == b'':
raise StopAsyncIteration
return val
if compat.PY352:
# В Python 3.5.2 и выше __aiter__ должен возвращать
# непосредственно асинхронный итератор.
def __aiter__(self):
return self