diff --git a/persistent_queue/persistent_queue.py b/persistent_queue/persistent_queue.py index 6634530..011fc86 100644 --- a/persistent_queue/persistent_queue.py +++ b/persistent_queue/persistent_queue.py @@ -131,8 +131,17 @@ def read_data(): timeout = target - time.time() elif not partial and self._length < items: + # If we are getting data and there is not enough data to get raise queue.Empty + elif partial and self._length == 0: + # We are peeking data (partial=True) and there is no data to get + # Save some disk IO by returning the result now + if items == 1: + raise queue.Empty + else: + return [], None + with self._file_lock: self._file.seek(self._get_queue_top(), 0) # Beginning of data total_items = self._length if items > self._length else items @@ -140,12 +149,9 @@ def read_data(): queue_top = self._file.tell() if items == 1: - if len(data) == 0: - _LOGGER.debug("No items to peek at so returning None") - return None, queue_top - else: - _LOGGER.debug("Returning data from peek") - return data[0], queue_top + _LOGGER.debug("Returning data from peek") + # We've insured that len(data) > 0 so data[0] is safe + return data[0], queue_top else: _LOGGER.debug("Returning data from peek") return data, queue_top @@ -283,9 +289,9 @@ def get(self, block=True, timeout=None, items=1): self._set_queue_top(queue_top) if isinstance(data, list): - if len(data) > 0: + if len(data) > 0: # Save a write to disk if we can self._update_length(self._length - len(data)) - elif data is not None: + else: self._update_length(self._length - 1) self._get_event.set() @@ -348,7 +354,10 @@ def peek(self, block=False, timeout=None, items=1): Peeks into the queue and returns items without removing them. """ with self._get_lock: - return self._peek(block, timeout, items, partial=True)[0] + try: + return self._peek(block, timeout, items, partial=True)[0] + except queue.Empty: + return None def clear(self): """ diff --git a/tests/test_persistent_queue.py b/tests/test_persistent_queue.py index 4e7c41d..3c7ecc3 100644 --- a/tests/test_persistent_queue.py +++ b/tests/test_persistent_queue.py @@ -262,6 +262,45 @@ def test_delete_no_values(self): self.queue.delete() self.queue.delete(100) + def test_push_none(self): + self.queue.put(1) + self.queue.put(None) + self.queue.put(None) + + assert len(self.queue) == 3 + assert self.queue.get() == 1 + + assert len(self.queue) == 2 + assert self.queue.get() is None + + assert len(self.queue) == 1 + assert self.queue.get() is None + + def test_return_types(self): + assert self.queue.peek(items=0) == [] + assert self.queue.peek(items=1) is None + assert self.queue.peek(items=2) == [] + assert self.queue.peek(items=100) == [] + + assert self.queue.get(items=0, block=False) == [] + with pytest.raises(queue.Empty): + self.queue.get(items=1, block=False) + with pytest.raises(queue.Empty): + self.queue.get(items=2, block=False) + with pytest.raises(queue.Empty): + self.queue.get(items=100, block=False) + + self.queue.put([1, 2, 3, 4, 5]) + + assert isinstance(self.queue.get(), int) + assert isinstance(self.queue.get(items=0), list) + assert isinstance(self.queue.get(items=1), int) + assert isinstance(self.queue.get(items=2), list) + + assert isinstance(self.queue.peek(items=0), list) + assert isinstance(self.queue.peek(items=1), int) + assert isinstance(self.queue.peek(items=2), list) + def test_big_file_1(self): data = {b'a': list(range(500))}