expose current_task_is_batch

Author: nsaje

tasktiger/tasktiger.py

@@ -317,13 +317,19 @@ def _get_current_serialized_func(self) -> str:
             raise RuntimeError("Must be accessed from within a task.")
         return g["current_tasks"][0].serialized_func
 
+    def _get_current_task_is_batch(self) -> bool:
+        if g["current_task_is_batch"] is None:
+            raise RuntimeError("Must be accessed from within a task.")
+        return g["current_task_is_batch"]
+
     """
     Properties to access the currently processing task (or tasks, in case of a
     batch task) from within the task. They must be invoked from within a task.
     """
     current_task = property(_get_current_task)
     current_tasks = property(_get_current_tasks)
     current_serialized_func = property(_get_current_serialized_func)
+    current_task_is_batch = property(_get_current_task_is_batch)
 
     @classproperty
     def current_instance(self) -> "TaskTiger":

tests/tasks.py

@@ -183,6 +183,23 @@ def verify_current_serialized_func_batch(tasks):
         conn.set("serialized_func", serialized_func)
 
 
+def verify_current_task_is_batch():
+    with redis.Redis(
+        host=REDIS_HOST, db=TEST_DB, decode_responses=True
+    ) as conn:
+        is_batch = tiger.current_task_is_batch
+        conn.set("current_task_is_batch", str(is_batch))
+
+
+@tiger.task(batch=True, queue="batch")
+def verify_current_task_is_batch_batch(tasks):
+    with redis.Redis(
+        host=REDIS_HOST, db=TEST_DB, decode_responses=True
+    ) as conn:
+        is_batch = tiger.current_task_is_batch
+        conn.set("current_task_is_batch", str(is_batch))
+
+
 @tiger.task()
 def verify_tasktiger_instance():
     # Not necessarily the same object, but the same configuration.

tests/test_base.py

@@ -49,6 +49,8 @@
     verify_current_serialized_func,
     verify_current_serialized_func_batch,
     verify_current_task,
+    verify_current_task_is_batch,
+    verify_current_task_is_batch_batch,
     verify_current_tasks,
     verify_tasktiger_instance,
 )
@@ -1200,6 +1202,37 @@ def test_current_serialized_func_batch(self, always_eager):
         )
 
 
+class TestCurrentTaskIsBatch(BaseTestCase):
+    """
+    Ensure current_task_is_batch is set.
+    """
+
+    @pytest.mark.parametrize("always_eager", [False, True])
+    def test_current_task_is_batch(self, always_eager):
+        self.tiger.config["ALWAYS_EAGER"] = always_eager
+        task = Task(self.tiger, verify_current_task_is_batch)
+        task.delay()
+        Worker(self.tiger).run(once=True)
+        assert not self.conn.exists("runtime_error")
+        assert (
+            self.conn.get("current_task_is_batch")
+            == "False"
+        )
+
+    @pytest.mark.parametrize("always_eager", [False, True])
+    def test_current_task_is_batch_batch(self, always_eager):
+        self.tiger.config["ALWAYS_EAGER"] = always_eager
+        task1 = Task(self.tiger, verify_current_task_is_batch_batch)
+        task1.delay()
+        task2 = Task(self.tiger, verify_current_task_is_batch_batch)
+        task2.delay()
+        Worker(self.tiger).run(once=True)
+        assert (
+            self.conn.get("current_task_is_batch")
+            == "True"
+        )
+
+
 class TestTaskTigerGlobal(BaseTestCase):
     """
     Ensure TaskTiger.current_instance is set.