Fix putting None in queue

persistent_queue/persistent_queue.py

@@ -131,21 +131,27 @@ 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
             data = [read_data() for i in range(total_items)]
             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):
         """

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))}