def __add__(self, other):
""" a + b means TaskGroup([a, b])
"""
if isiterable(other) and not isinstance(other, AbstractDAG):
return TaskGroup([self] + list(other))
else:
return TaskGroup((self, other))
def _add_edge(self, task_from, task_to, group_name=None):
"""Add an edge between two tasks
Parameters
----------
group_name : str
Pass a string to group this edge, upon rendering, upstream
products are available via task[group_name][tas_name]
"""
attrs = {} if group_name is None else {'group_name': group_name}
# when adding a task group (but not a dag)
if isiterable(task_from) and not isinstance(task_from, DAG):
# if iterable, add all components as separate upstream tasks
for a_task_from in task_from:
# this happens when the task was originally declared in
# another dag...
if a_task_from.name not in self._G:
self._G.add_node(a_task_from.name, task=a_task_from)
self._G.add_edge(a_task_from.name, task_to.name, **attrs)
else:
# this happens when the task was originally declared in
# another dag...
if task_from.name not in self._G:
self._G.add_node(task_from.name, task=task_from)
# DAGs are treated like a single task
self._G.add_edge(task_from.name, task_to.name, **attrs)
def _add_edge(self, task_from, task_to):
"""Add an edge between two tasks
"""
# if a new task is added, rendering is required again
self._did_render = False
if isiterable(task_from) and not isinstance(task_from, DAG):
# if iterable, add all components as separate upstream tasks
for a_task_from in task_from:
# this happens when the task was originally declared in
# another dag...
if a_task_from.name not in self._G:
self._G.add_node(a_task_from.name, task=a_task_from)
self._G.add_edge(a_task_from.name, task_to.name)
else:
# this happens when the task was originally declared in
# another dag...
if task_from.name not in self._G:
self._G.add_node(task_from.name, task=task_from)
# DAGs are treated like a single task
self._G.add_edge(task_from.name, task_to.name)
def set_upstream(self, other):
if isiterable(other):
for t in self.tasks:
for o in other:
t.set_upstream(other)
else:
for t in self.tasks:
t.set_upstream(other)
def set_upstream(self, other):
# if self.treat_as_single_task:
# if isiterable(other):
# for o in other:
# self.dag._add_edge(o, self)
# else:
# self.dag._add_edge(other, self)
# else:
if isiterable(other):
for t in self.tasks:
for o in other:
t.set_upstream(other)
else:
for t in self.tasks:
t.set_upstream(other)
def __radd__(self, other):
if isiterable(other):
return TaskGroup(list(other) + list(self.tasks))
else:
return TaskGroup([other] + list(self.tasks))