refactor: Add _add_provenance method to TaskWorker class

provos · Sep 12, 2024 · d6d4648 · d6d4648
1 parent cc2dd8d
commit d6d4648
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/src/planai/task.py b/src/planai/task.py
@@ -370,6 +370,12 @@ def consume_work(self, task: Task):
         """
         pass
 
+    def _add_provenance(self, task: Task) -> Task:
+        with self._state_lock:
+            self._id += 1
+            task._provenance.append((self.name, self._id))
+        return task
+
     def publish_work(self, task: Task, input_task: Optional[Task]):
         """
         Publish a work item.
@@ -398,10 +404,8 @@ def publish_work(self, task: Task, input_task: Optional[Task]):
             task._provenance = []
             task._input_provenance = []
 
-        with self._state_lock:
-            self._id += 1
+        self._add_provenance(task)
 
-        task._provenance.append((self.name, self._id))
         logging.info(
             "Task %s published work with provenance %s", self.name, task._provenance
         )