Get sizes for queues and states (#352)

Lets us query specific queue sizes in a single Redis round trip.

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## Version 0.20
+
+* Add `tiger.get_sizes_for_queues_and_states` ([352](https://github.com/closeio/tasktiger/pull/352))
+
 ## Version 0.19.5
 
 * First version using the automated-release process

tasktiger/__init__.py

@@ -12,7 +12,7 @@
 from .tasktiger import TaskTiger, run_worker
 from .worker import Worker
 
-__version__ = "0.19.5"
+__version__ = "0.20"
 __all__ = [
     "TaskTiger",
     "Worker",

tasktiger/tasktiger.py

@@ -489,6 +489,23 @@ def get_queue_sizes(self, queue: str) -> Dict[str, int]:
         results = pipeline.execute()
         return dict(zip(states, results))
 
+    def get_sizes_for_queues_and_states(
+        self, queues_and_states: List[Tuple[str, str]]
+    ) -> List[int]:
+        """
+        Get the sizes for the specific queues and states.
+
+        queues_and_states: List of tuples (queue_name, state).
+
+        Returns a list of queue sizes in the order of the passed
+        queues_and_states.
+        """
+        pipeline = self.connection.pipeline()
+        for queue, state in queues_and_states:
+            pipeline.zcard(self._key(state, queue))
+        results = pipeline.execute()
+        return results
+
     def get_total_queue_size(self, queue: str) -> int:
         """Get total queue size for QUEUED, SCHEDULED, and ACTIVE states."""
 

tests/test_queue_size.py

@@ -88,3 +88,38 @@ def test_task_all_states(self):
                 max_queue_size=3,
                 when=datetime.timedelta(seconds=10),
             )
+
+
+class TestQueueSizes:
+    @pytest.fixture
+    def queue_sample_tasks(self, tiger):
+        tiger.delay(simple_task)
+        tiger.delay(simple_task)
+        tiger.delay(simple_task, queue="other")
+        tiger.delay(simple_task, when=datetime.timedelta(seconds=60))
+
+    def test_get_total_queue_size(self, tiger, queue_sample_tasks):
+        assert tiger.get_total_queue_size("other") == 1
+        assert tiger.get_total_queue_size("default") == 3
+
+    def test_get_queue_sizes(self, tiger, queue_sample_tasks):
+        assert tiger.get_queue_sizes("default") == {
+            "active": 0,
+            "queued": 2,
+            "scheduled": 1,
+        }
+        assert tiger.get_queue_sizes("other") == {
+            "active": 0,
+            "queued": 1,
+            "scheduled": 0,
+        }
+
+    def test_get_sizes_for_queues_and_states(self, tiger, queue_sample_tasks):
+        assert tiger.get_sizes_for_queues_and_states(
+            [
+                ("default", "queued"),
+                ("default", "scheduled"),
+                ("other", "queued"),
+                ("other", "scheduled"),
+            ]
+        ) == [2, 1, 1, 0]