queues.py
python
__all__ = (
'Queue',
'PriorityQueue',
'LifoQueue',
'QueueFull',
'QueueEmpty',
'QueueShutDown',
)
import collections
import heapq
from types import GenericAlias
from . import locks
from . import mixins
class QueueEmpty(Exception):
"""Вызывается, когда Queue.get_nowait() вызывается на пустой Queue."""
pass
class QueueFull(Exception):
"""Вызывается, когда метод Queue.put_nowait() вызывается на полной Queue."""
pass
class QueueShutDown(Exception):
"""Возникает при помещении в или извлечении из закрытой очереди."""
pass
class Queue(mixins._LoopBoundMixin):
"""Очередь, полезная для координации корутин-производителя и корутин-потребителя.
Если maxsize меньше или равен нулю, размер очереди бесконечен.
Если это целое число больше 0, то "await put()" будет блокироваться, когда
очередь достигает maxsize, до тех пор, пока элемент не будет удалён вызовом get().
В отличие от queue.Queue, вы можете надёжно узнать размер этой очереди
с помощью qsize(), так как ваше однопоточное приложение asyncio не будет
прервано между вызовом qsize() и выполнением операции над очередью.
"""
def __init__(self, maxsize=0):
self._maxsize = maxsize
# Future.
self._getters = collections.deque()
# Future.
self._putters = collections.deque()
self._unfinished_tasks = 0
self._finished = locks.Event()
self._finished.set()
self._init(maxsize)
self._is_shutdown = False
# Эти три метода можно переопределить в подклассах.
def _init(self, maxsize):
self._queue = collections.deque()
def _get(self):
return self._queue.popleft()
def _put(self, item):
self._queue.append(item)
# Конец переопределяемых методов.
def _wakeup_next(self, waiters):
# Разбудить следующего ожидающего (если есть), который не отменен.
while waiters:
waiter = waiters.popleft()
if not waiter.done():
waiter.set_result(None)
break
def __repr__(self):
return f'<{type(self).__name__} at {id(self):#x} {self._format()}>'
def __str__(self):
return f'<{type(self).__name__} {self._format()}>'
__class_getitem__ = classmethod(GenericAlias)
def _format(self):
result = f'maxsize={self._maxsize!r}'
if getattr(self, '_queue', None):
result += f' _queue={list(self._queue)!r}'
if self._getters:
result += f' _getters[{len(self._getters)}]'
if self._putters:
result += f' _putters[{len(self._putters)}]'
if self._unfinished_tasks:
result += f' tasks={self._unfinished_tasks}'
if self._is_shutdown:
result += ' shutdown'
return result
def qsize(self):
"""Количество элементов в очереди."""
return len(self._queue)
@property
def maxsize(self):
"""Максимальное количество элементов в очереди."""
return self._maxsize
def empty(self):
"""Возвращает True, если очередь пуста, иначе False."""
return not self._queue
def full(self):
"""Возвращает True, если в очереди maxsize элементов.
Примечание: если Queue была инициализирована с maxsize=0 (по умолчанию),
то full() никогда не возвращает True.
"""
if self._maxsize <= 0:
return False
else:
return self.qsize() >= self._maxsize
async def put(self, item):
"""Поместить элемент в очередь.
Поместить элемент в очередь. Если очередь полна, ждать, пока не появится
свободное место, прежде чем добавить элемент.
Вызывает QueueShutDown, если очередь была закрыта.
"""
while self.full():
if self._is_shutdown:
raise QueueShutDown
putter = self._get_loop().create_future()
self._putters.append(putter)
try:
await putter
except:
putter.cancel() # На случай, если putter еще не завершен.
try:
# Очищает self._putters от отмененных putter'ов.
self._putters.remove(putter)
except ValueError:
# Putter может быть удален из self._putters с помощью
# предыдущий вызов get_nowait или вызов shutdown.
pass
if not self.full() and not putter.cancelled():
# Мы были разбужены get_nowait(), но не можем принять
# вызов. Пробуждаем следующего в очереди.
self._wakeup_next(self._putters)
raise
return self.put_nowait(item)
def put_nowait(self, item):
"""Поместить элемент в очередь без блокировки.
Если свободное место недоступно немедленно, вызвать QueueFull.
Вызывает QueueShutDown, если очередь была закрыта.
"""
if self._is_shutdown:
raise QueueShutDown
if self.full():
raise QueueFull
self._put(item)
self._unfinished_tasks += 1
self._finished.clear()
self._wakeup_next(self._getters)
async def get(self):
"""Удалить и вернуть элемент из очереди.
Если очередь пуста, ждать, пока элемент не станет доступен.
Вызывает QueueShutDown, если очередь была закрыта и пуста,
или если очередь была закрыта немедленно.
"""
while self.empty():
if self._is_shutdown and self.empty():
raise QueueShutDown
getter = self._get_loop().create_future()
self._getters.append(getter)
try:
await getter
except:
getter.cancel() # На случай, если getter еще не завершен.
try:
# Очищает self._getters от отмененных getter'ов.
self._getters.remove(getter)
except ValueError:
# Getter может быть удален из self._getters с помощью
# предыдущий вызов put_nowait или вызов shutdown.
pass
if not self.empty() and not getter.cancelled():
# Мы были пробуждены с помощью put_nowait(), но не можем взять
# вызов. Пробуждаем следующего в очереди.
self._wakeup_next(self._getters)
raise
return self.get_nowait()
def get_nowait(self):
"""Удалить и вернуть элемент из очереди.
Возвращает элемент, если он немедленно доступен, иначе вызывает
QueueEmpty.
Вызывает QueueShutDown, если очередь была закрыта и пуста,
или если очередь была закрыта немедленно.
"""
if self.empty():
if self._is_shutdown:
raise QueueShutDown
raise QueueEmpty
item = self._get()
self._wakeup_next(self._putters)
return item
def task_done(self):
"""Указать, что ранее поставленная задача завершена.
Используется потребителями очереди. Для каждого get(), использованного для получения задачи,
последующий вызов task_done() сообщает очереди, что обработка задачи завершена.
Если join() в данный момент блокируется, он возобновится, когда все элементы
будут обработаны (то есть для каждого элемента, помещенного в очередь с помощью put(),
был получен вызов task_done()).
Вызывает ValueError, если вызван больше раз, чем было помещено элементов
в очередь.
"""
if self._unfinished_tasks <= 0:
raise ValueError('task_done() called too many times')
self._unfinished_tasks -= 1
if self._unfinished_tasks == 0:
self._finished.set()
async def join(self):
"""Блокировать, пока все элементы в очереди не будут получены и обработаны.
Счетчик незавершенных задач увеличивается при добавлении элемента в
очередь. Счетчик уменьшается, когда потребитель вызывает
task_done() для указания того, что элемент был получен и вся работа над
ним завершена. Когда счетчик незавершенных задач падает до нуля,
join() разблокируется.
"""
if self._unfinished_tasks > 0:
await self._finished.wait()
def shutdown(self, immediate=False):
"""Завершить работу очереди, заставляя get и put вызывать QueueShutDown.
По умолчанию get будут вызывать исключение только когда очередь пуста. Установите
'immediate' в True, чтобы get вызывали исключение немедленно.
Все заблокированные вызывающие put() и get() будут разблокированы.
Если 'immediate', очередь очищается, и количество незавершенных задач
уменьшается на количество очищенных задач. Если количество незавершенных задач
уменьшается до нуля, вызывающие Queue.join разблокируются.
"""
self._is_shutdown = True
if immediate:
while not self.empty():
self._get()
if self._unfinished_tasks > 0:
self._unfinished_tasks -= 1
if self._unfinished_tasks == 0:
self._finished.set()
# Все получатели должны повторно проверить пустоту очереди, чтобы вызвать ShutDown
while self._getters:
getter = self._getters.popleft()
if not getter.done():
getter.set_result(None)
while self._putters:
putter = self._putters.popleft()
if not putter.done():
putter.set_result(None)
class PriorityQueue(Queue):
"""Подкласс Queue; извлекает элементы в порядке приоритета (сначала наименьший).
Элементы обычно представлены кортежами вида: (номер приоритета, данные).
"""
def _init(self, maxsize):
self._queue = []
def _put(self, item, heappush=heapq.heappush):
heappush(self._queue, item)
def _get(self, heappop=heapq.heappop):
return heappop(self._queue)
class LifoQueue(Queue):
"""Подкласс Queue, который первым извлекает самые последние добавленные элементы."""
def _init(self, maxsize):
self._queue = []
def _put(self, item):
self._queue.append(item)
def _get(self):
return self._queue.pop()