Add helper to workflows to get chunks of branches.
@@ -22,7 +22,7 @@ from law.task.base import Task, Register
from law.task.proxy import ProxyTask, get_proxy_attribute
from law.target.collection import TargetCollection
from law.parameter import NO_STR, NO_INT, CSVParameter
-from law.util import no_value, make_list, DotDict
+from law.util import no_value, make_list, iter_chunks, DotDict
logger = logging.getLogger(__name__)
@@ -607,7 +607,7 @@ class BaseWorkflow(Task):
"""
if self.is_branch():
return self.as_workflow().get_branch_tasks()
- else:
+
if self._branch_tasks is None:
branch_map = self.get_branch_map()
if branch_map is None:
@@ -620,6 +620,37 @@ class BaseWorkflow(Task):
return self._branch_tasks
+ def get_branch_chunks(self, chunk_size, new_inst=True, **kwargs):
+ """
+ Returns a list of chunks of branch numbers with a given size *chunk_size*. When *new_inst*
+ is *True*, a new workflow instance is created using :py:meth:`BaseTask.req`, forwarding all
+ *kwargs*. Its *_exclude* list will contain ``["start_branch", "end_branch", "branches"]`` in
+ order to use all possible branch values. Example:
+
+ .. code-block:: python
+
+ wf = SomeWorkflowTask() # has 8 branches
+
+ print(wf.get_branch_chunks(3))
+ # -> [[0, 1, 2], [3, 4, 5], [6, 7]]
+ """
+ if self.is_branch():
+ return self.as_workflow().get_branch_chunks(chunk_size, new_inst=new_inst, **kwargs)
+
+ # create a new instance if required
+ if new_inst:
+ _exclude = make_list(kwargs.get("_exclude", []))
+ _exclude.extend(["start_branch", "end_branch", "branches"])
+ kwargs["_exclude"] = _exclude
+ inst = self.req(self, **kwargs)
+ else:
+ inst = self
+
+ # get the branch map and create chunks of its branch values
+ branch_chunks = iter_chunks(inst.branch_map.keys(), chunk_size)
+
+ return list(branch_chunks)
+
def workflow_requires(self):
"""
Hook to add workflow requirements. This method is expected to return a dictionary. When
