diff --git a/src/textual/message_pump.py b/src/textual/message_pump.py index 8d80d8bab1..b869d9bd5c 100644 --- a/src/textual/message_pump.py +++ b/src/textual/message_pump.py @@ -14,7 +14,6 @@ import threading from abc import ABC, abstractmethod from asyncio import CancelledError, QueueEmpty, Task, create_task -from collections import deque from contextlib import contextmanager from functools import partial from time import perf_counter @@ -25,7 +24,6 @@ Callable, ClassVar, Generator, - Generic, Iterable, Type, TypeVar, @@ -105,55 +103,26 @@ def call_after(self, pump: MessagePump, callback: Callable[[], object]) -> bool: QueueType = TypeVar("QueueType") -class Queue(Generic[QueueType]): - """A cut-down asyncio.Queue: just enough to run a message pump, and whether it owes work.""" +class Queue(asyncio.Queue[QueueType]): + """asyncio's queue, plus a peek at its head and whether the pump owes work.""" - def __init__(self) -> None: - self.values: deque[QueueType] = deque() - self.ready_event = asyncio.Event() - self.activity: QueueActivity = SETTLED + activity: QueueActivity = SETTLED - def put_nowait(self, value: QueueType) -> None: - self.values.append(value) + def _put(self, item: QueueType) -> None: + super()._put(item) self.activity = OWING - self.ready_event.set() - - def qsize(self) -> int: - return len(self.values) - - def empty(self) -> bool: - return not self.values - - def task_done(self) -> None: - pass async def get(self) -> QueueType: - # An Event wakes every waiter; a synchronous consumer may also drain - # the item before an awakened task resumes. Notification is not a - # reservation, so recheck the queue after each wait. - while not self.values: + if self.empty(): # Asking with none left: every value put before has been handled. self.activity = SETTLED - self.ready_event.clear() - await self.ready_event.wait() - value = self.values.popleft() - if not self.values: - self.ready_event.clear() - return value + return await super().get() def peek_nowait(self) -> QueueType: """Read the next value without taking it from a waiting consumer.""" - if not self.values: - raise QueueEmpty() - return self.values[0] - - def get_nowait(self) -> QueueType: - if not self.values: + if self.empty(): raise QueueEmpty() - value = self.values.popleft() - if not self.values: - self.ready_event.clear() - return value + return self._queue[0] _MessagePumpMetaSub = TypeVar("_MessagePumpMetaSub", bound="_MessagePumpMeta") diff --git a/tests/test_queue.py b/tests/test_queue.py index f58b87c74c..2264dafd1c 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -2,59 +2,36 @@ import pytest -from textual.message_pump import Queue +from textual.message_pump import OWING, SETTLED, Queue -async def test_get_rechecks_after_nowait_consumer_drains_wakeup(): +async def test_peek_leaves_the_head_for_a_waiting_consumer(): queue = Queue() waiting = asyncio.create_task(queue.get()) try: await asyncio.sleep(0) - queue.put_nowait("first") - assert queue.get_nowait() == "first" - await asyncio.sleep(0) - assert not waiting.done() - queue.put_nowait("second") - assert await asyncio.wait_for(waiting, 1) == "second" + with pytest.raises(asyncio.QueueEmpty): + queue.peek_nowait() + queue.put_nowait("head") + assert queue.peek_nowait() == "head" + assert await asyncio.wait_for(waiting, 1) == "head" assert queue.empty() finally: waiting.cancel() await asyncio.gather(waiting, return_exceptions=True) -async def test_one_item_does_not_satisfy_two_waiting_consumers(): +async def test_owing_from_a_put_until_a_get_finds_none_left(): queue = Queue() - first = asyncio.create_task(queue.get()) - second = asyncio.create_task(queue.get()) - try: - await asyncio.sleep(0) - queue.put_nowait("first") - assert await asyncio.wait_for(first, 1) == "first" - assert not second.done() - queue.put_nowait("second") - assert await asyncio.wait_for(second, 1) == "second" - assert queue.empty() - finally: - first.cancel() - second.cancel() - await asyncio.gather(first, second, return_exceptions=True) - - -async def test_cancelled_waiter_does_not_lose_wakeup(): - queue = Queue() - first = asyncio.create_task(queue.get()) - second = asyncio.create_task(queue.get()) + assert queue.activity is SETTLED + queue.put_nowait("message") + assert queue.activity is OWING + assert await queue.get() == "message" + assert queue.activity is OWING + waiting = asyncio.create_task(queue.get()) try: await asyncio.sleep(0) - queue.put_nowait(None) - first.cancel() - with pytest.raises(asyncio.CancelledError): - await first - assert await asyncio.wait_for(second, 1) is None - with pytest.raises(asyncio.QueueEmpty): - queue.get_nowait() - assert not queue.ready_event.is_set() + assert queue.activity is SETTLED finally: - first.cancel() - second.cancel() - await asyncio.gather(first, second, return_exceptions=True) + waiting.cancel() + await asyncio.gather(waiting, return_exceptions=True)