1__all__ = ('Queue', 'PriorityQueue', 'LifoQueue', 'QueueFull', 'QueueEmpty')
2
3import collections
4import heapq
5
6from . import locks
7from . import mixins
8
9
10class QueueEmpty(Exception):
11    """Raised when Queue.get_nowait() is called on an empty Queue."""
12    pass
13
14
15class QueueFull(Exception):
16    """Raised when the Queue.put_nowait() method is called on a full Queue."""
17    pass
18
19
20class Queue(mixins._LoopBoundMixin):
21    """A queue, useful for coordinating producer and consumer coroutines.
22
23    If maxsize is less than or equal to zero, the queue size is infinite. If it
24    is an integer greater than 0, then "await put()" will block when the
25    queue reaches maxsize, until an item is removed by get().
26
27    Unlike the standard library Queue, you can reliably know this Queue's size
28    with qsize(), since your single-threaded asyncio application won't be
29    interrupted between calling qsize() and doing an operation on the Queue.
30    """
31
32    def __init__(self, maxsize=0, *, loop=mixins._marker):
33        super().__init__(loop=loop)
34        self._maxsize = maxsize
35
36        # Futures.
37        self._getters = collections.deque()
38        # Futures.
39        self._putters = collections.deque()
40        self._unfinished_tasks = 0
41        self._finished = locks.Event()
42        self._finished.set()
43        self._init(maxsize)
44
45    # These three are overridable in subclasses.
46
47    def _init(self, maxsize):
48        self._queue = collections.deque()
49
50    def _get(self):
51        return self._queue.popleft()
52
53    def _put(self, item):
54        self._queue.append(item)
55
56    # End of the overridable methods.
57
58    def _wakeup_next(self, waiters):
59        # Wake up the next waiter (if any) that isn't cancelled.
60        while waiters:
61            waiter = waiters.popleft()
62            if not waiter.done():
63                waiter.set_result(None)
64                break
65
66    def __repr__(self):
67        return f'<{type(self).__name__} at {id(self):#x} {self._format()}>'
68
69    def __str__(self):
70        return f'<{type(self).__name__} {self._format()}>'
71
72    def __class_getitem__(cls, type):
73        return cls
74
75    def _format(self):
76        result = f'maxsize={self._maxsize!r}'
77        if getattr(self, '_queue', None):
78            result += f' _queue={list(self._queue)!r}'
79        if self._getters:
80            result += f' _getters[{len(self._getters)}]'
81        if self._putters:
82            result += f' _putters[{len(self._putters)}]'
83        if self._unfinished_tasks:
84            result += f' tasks={self._unfinished_tasks}'
85        return result
86
87    def qsize(self):
88        """Number of items in the queue."""
89        return len(self._queue)
90
91    @property
92    def maxsize(self):
93        """Number of items allowed in the queue."""
94        return self._maxsize
95
96    def empty(self):
97        """Return True if the queue is empty, False otherwise."""
98        return not self._queue
99
100    def full(self):
101        """Return True if there are maxsize items in the queue.
102
103        Note: if the Queue was initialized with maxsize=0 (the default),
104        then full() is never True.
105        """
106        if self._maxsize <= 0:
107            return False
108        else:
109            return self.qsize() >= self._maxsize
110
111    async def put(self, item):
112        """Put an item into the queue.
113
114        Put an item into the queue. If the queue is full, wait until a free
115        slot is available before adding item.
116        """
117        while self.full():
118            putter = self._get_loop().create_future()
119            self._putters.append(putter)
120            try:
121                await putter
122            except:
123                putter.cancel()  # Just in case putter is not done yet.
124                try:
125                    # Clean self._putters from canceled putters.
126                    self._putters.remove(putter)
127                except ValueError:
128                    # The putter could be removed from self._putters by a
129                    # previous get_nowait call.
130                    pass
131                if not self.full() and not putter.cancelled():
132                    # We were woken up by get_nowait(), but can't take
133                    # the call.  Wake up the next in line.
134                    self._wakeup_next(self._putters)
135                raise
136        return self.put_nowait(item)
137
138    def put_nowait(self, item):
139        """Put an item into the queue without blocking.
140
141        If no free slot is immediately available, raise QueueFull.
142        """
143        if self.full():
144            raise QueueFull
145        self._put(item)
146        self._unfinished_tasks += 1
147        self._finished.clear()
148        self._wakeup_next(self._getters)
149
150    async def get(self):
151        """Remove and return an item from the queue.
152
153        If queue is empty, wait until an item is available.
154        """
155        while self.empty():
156            getter = self._get_loop().create_future()
157            self._getters.append(getter)
158            try:
159                await getter
160            except:
161                getter.cancel()  # Just in case getter is not done yet.
162                try:
163                    # Clean self._getters from canceled getters.
164                    self._getters.remove(getter)
165                except ValueError:
166                    # The getter could be removed from self._getters by a
167                    # previous put_nowait call.
168                    pass
169                if not self.empty() and not getter.cancelled():
170                    # We were woken up by put_nowait(), but can't take
171                    # the call.  Wake up the next in line.
172                    self._wakeup_next(self._getters)
173                raise
174        return self.get_nowait()
175
176    def get_nowait(self):
177        """Remove and return an item from the queue.
178
179        Return an item if one is immediately available, else raise QueueEmpty.
180        """
181        if self.empty():
182            raise QueueEmpty
183        item = self._get()
184        self._wakeup_next(self._putters)
185        return item
186
187    def task_done(self):
188        """Indicate that a formerly enqueued task is complete.
189
190        Used by queue consumers. For each get() used to fetch a task,
191        a subsequent call to task_done() tells the queue that the processing
192        on the task is complete.
193
194        If a join() is currently blocking, it will resume when all items have
195        been processed (meaning that a task_done() call was received for every
196        item that had been put() into the queue).
197
198        Raises ValueError if called more times than there were items placed in
199        the queue.
200        """
201        if self._unfinished_tasks <= 0:
202            raise ValueError('task_done() called too many times')
203        self._unfinished_tasks -= 1
204        if self._unfinished_tasks == 0:
205            self._finished.set()
206
207    async def join(self):
208        """Block until all items in the queue have been gotten and processed.
209
210        The count of unfinished tasks goes up whenever an item is added to the
211        queue. The count goes down whenever a consumer calls task_done() to
212        indicate that the item was retrieved and all work on it is complete.
213        When the count of unfinished tasks drops to zero, join() unblocks.
214        """
215        if self._unfinished_tasks > 0:
216            await self._finished.wait()
217
218
219class PriorityQueue(Queue):
220    """A subclass of Queue; retrieves entries in priority order (lowest first).
221
222    Entries are typically tuples of the form: (priority number, data).
223    """
224
225    def _init(self, maxsize):
226        self._queue = []
227
228    def _put(self, item, heappush=heapq.heappush):
229        heappush(self._queue, item)
230
231    def _get(self, heappop=heapq.heappop):
232        return heappop(self._queue)
233
234
235class LifoQueue(Queue):
236    """A subclass of Queue that retrieves most recently added entries first."""
237
238    def _init(self, maxsize):
239        self._queue = []
240
241    def _put(self, item):
242        self._queue.append(item)
243
244    def _get(self):
245        return self._queue.pop()
246