'''
A tagged queue. The producer(s) place items in the queue, each item _may_ have
a list of tags attached to it. Consumer(s) may request an item of a given tag.
The tag is then "consumed" from this item. The item is removed from the queue
when all its tags are consumed.
For better understanding, think about it as set of discinct queues,
one for a given tag (and one for no tag).
This will be appropriate for most cases (although internal
plumbing is different), except for full() which counts all
items regardless of tags to guarrantee maximum object count in memory
If tags are not used, the TaggedQueue behaves exactly as Queue.Queue,
but performance is worse (see below).
Assumptions:
- tag can be anything that works as a dictionary index
- when putting an item in the queue, any number of tags can be used,
but list must be provided (or the tag_list parameter not used at all)
- when getting an item for the queue, up to 1 tag can be provided.
Example:
>>> q.put('red only', tag_list=['red'])
>>> q.put('red and blue', tag_list=['red', 'blue'])
>>> q.put('blue only', tag_list=['blue'])
>>> q.put('both colours', tag_list=['red', 'blue'])
>>> q.get(tag='blue')
'red and blue'
>>> q.get(tag='blue')
'blue only'
>>> q.get(tag='red')
'red only'
>>> q.get(tag='red')
'red and blue'
>>> q.get(tag='red')
'both colours'
>>> q.get(tag='blue')
'both colours'
Performance (in seconds, smaller is better):
- adding an object in the queue 100000 times, then getting it back 100000 times
(max queue length is 100000)
Queue.Queue: 1.124342
TaggedQueue without tags: 5.291548
TaggedQueue with 10 distinct tags: 2.100272
TaggedQueue with 1000 distinct tags: 1.692542
- 100000 consecutive put() - get() calls, (max queue length is 1)
Queue.Queue: 1.158024
TaggedQueue without tags: 2.218153
TaggedQueue with 10 distinct tags: 2.270345
TaggedQueue with 1000 distinct tags: 2.344355
'''
import multiprocessing
from sys import maxint
from time import time as _time
def id_generator():
'''a generator returning numeric ID's as identifiers'''
_id = 0
while True:
_id += 1
if _id == maxint:
_id = 0
yield _id
class Empty(Exception):
'''raised when TaggedQueue is empty for a given tag'''
class Full(Exception):
'''raised when TaggedQueue is full (regardless of tags)'''
NOTAG = '_notag'
# based on Queue.Queue (no surprise)
class TaggedQueue(object):
'''A queue similar to Queue.Queue, but producers can tag items
and consumers can specify what tag are they interested in.
If an item has multiple tags, it is removed from the queue only after
all tags attached to it have been consumed.
'''
def __init__(self, maxsize=0):
'''initialize the queue'''
self.mutex = multiprocessing.Lock()
self.not_empty = multiprocessing.Condition(self.mutex)
self.not_full = multiprocessing.Condition(self.mutex)
self.maxsize = maxsize
self._tags = {} # list of refid's for each tag
self._queue = {} # the actual queue data
self._refcount = {} # how many tags refer to a given refid in the queue
self.id_generator = id_generator()
def qsize(self, tag=NOTAG):
'''Return the approximate size of the queue for a given tag.'''
self.mutex.acquire()
try:
n = len(self._tags[tag])
except KeyError:
n = 0
self.mutex.release()
return n
def empty(self, tag=NOTAG):
'''Return True if the queue is empty for a given tag'''
self.mutex.acquire()
n = self._empty(tag)
self.mutex.release()
return n
def full(self):
'''Return True if the queue is full, regardless of tags'''
self.mutex.acquire()
n = self._full()
self.mutex.release()
return n
def get(self, tag=NOTAG, block=True, timeout=None):
'''get item object with defined tag from the queue, assuming only
one tag
'''
self.not_empty.acquire()
try:
if not block:
if self._empty(tag):
raise Empty
elif timeout is None:
while self._empty(tag):
self.not_empty.wait()
else:
if timeout < 0:
raise ValueError("'timeout' must be a positive number")
endtime = _time() + timeout
while self._empty(tag):
remaining = endtime - _time()
if remaining <= 0.0:
raise Empty
self.not_empty.wait(remaining)
item = self._get(tag)
self.not_full.notify()
return item
finally:
self.not_empty.release()
def get_nowait(self, tag=NOTAG):
'''get tagged item immediatelly or raise an exception'''
return self.get(tag, block=False)
def put(self, item, tag_list=[NOTAG], block=True, timeout=None):
'''put an item in queue, tag_list is an optional list of tags
to be attached to the item
'''
self.not_full.acquire()
try:
if not block:
if self._full():
raise Full
elif timeout is None:
while self._full():
self.not_full.wait()
else:
if timeout < 0:
raise ValueError("'timeout' must be a positive number")
endtime = _time() + timeout
while self._full():
remaining = endtime - _time()
if remaining <= 0.0:
raise Full
self.not_full.wait(remaining)
self._put(item, tag_list)
self.not_empty.notify()
finally:
self.not_full.release()
def put_nowait(self, item, tag_list=[NOTAG]):
'''put an item to the list immediatelly or raise an exception'''
return self.put(item, tag_list, False)
def _empty(self, tag):
return not (tag in self._tags)
def _full(self):
return self.maxsize > 0 and len(self._refcount) == self.maxsize
def _get(self, tag):
if not tag in self._tags:
raise Empty
# get item's refid of the item
refid = self._tags[tag].pop(0) # get and remove first item from list
if not self._tags[tag]: # no item in the queue have this tag anymore
del(self._tags[tag])
# get the item from the queue
item = self._queue[refid]
# decrease the reference counter
self._refcount[refid] -= 1
if self._refcount[refid] == 0: # no more references
del self._refcount[refid]
del self._queue[refid]
return item
def _put(self, item, tag_list):
refid = self.id_generator.next()
self._queue[refid] = item
self._refcount[refid] = len(tag_list)
for t in tag_list:
try:
self._tags[t].append(refid)
except KeyError:
self._tags[t] = [refid]
## --- test_tagged_queue.py
'''
Test that TaggedQueue works as expected
'''
from tagged_queue import TaggedQueue, Empty
def test_empty_queue():
'''test that TaggedQueue returns Empty correctly'''
q = TaggedQueue()
try:
q.get_nowait()
assert False
except Empty:
pass
try:
q.get_nowait(tag='non_existing')
assert False
except Empty:
pass
def test_put_and_get():
'''test that tags are honoured for basic queue operations'''
q = TaggedQueue()
q.put_nowait('no tag')
assert q.get_nowait() == 'no tag'
q.put_nowait('test data', tag_list=['tag1', 'tag2', 3])
try:
q.get_nowait()
assert False
except Empty:
pass
q.get_nowait(tag=3)
q.get_nowait(tag='tag1')
q.get_nowait(tag='tag2')
try:
q.get_nowait(tag='tag2')
assert False
except Empty:
pass
def test_empty():
'''test that .empty() function works as expected'''
q = TaggedQueue()
assert q.empty()
assert q.empty(tag='tag1')
q.put_nowait('test data')
assert not q.empty()
assert q.empty(tag='tag1')
q.put_nowait('other test data', tag_list=['tag1'])
assert not q.empty()
assert not q.empty(tag='tag1')
q.get_nowait()
assert q.empty()
assert not q.empty(tag='tag1')
q.get_nowait(tag='tag1')
assert q.empty()
assert q.empty('tag1')
def test_qsize():
'''test that .qsize() function works as expected'''
q = TaggedQueue()
assert q.qsize() == 0
assert q.qsize(tag='tag1') == 0
q.put_nowait('test data')
assert q.qsize() == 1
assert q.qsize(tag='tag1') == 0
q.put_nowait('other test data', tag_list=['tag1'])
assert q.qsize() == 1
assert q.qsize(tag='tag1') == 1
q.put_nowait('other test data', tag_list=['tag1'])
assert q.qsize() == 1
assert q.qsize(tag='tag1') == 2
q.get_nowait()
assert q.qsize() == 0
assert q.qsize(tag='tag1') == 2
q.get_nowait(tag='tag1')
assert q.qsize() == 0
assert q.qsize(tag='tag1') == 1
q.get_nowait(tag='tag1')
assert q.qsize() == 0
assert q.qsize(tag='tag1') == 0
def test_full():
'''test that .full() function works as expected'''
q = TaggedQueue(maxsize=3)
assert not q.full()
q.put_nowait('data')
assert not q.full()
q.put_nowait('other data', tag_list=['some_tag'])
assert not q.full()
q.put_nowait('yet another data')
assert q.full()
def test_ordering():
'''test that queue ordering works properly'''
q = TaggedQueue()
q.put_nowait('red only', tag_list=['red'])
q.put_nowait('red and blue', tag_list=['red', 'blue'])
q.put_nowait('blue only', tag_list=['blue'])
q.put_nowait('both colours', tag_list=['red', 'blue'])
assert q.get_nowait(tag='blue') == 'red and blue'
assert q.get_nowait(tag='blue') == 'blue only'
assert q.get_nowait(tag='red') == 'red only'
assert q.get_nowait(tag='red') == 'red and blue'
assert q.get_nowait(tag='red') == 'both colours'
assert q.get_nowait(tag='blue') == 'both colours'
