def _build(self, *, data_source=None, **kwargs):
local_flow = LocalDatasetFlow(**kwargs)
qtzal_flow = QuetzalDatasetFlow(**kwargs)
self.update(local_flow)
self.update(qtzal_flow)
merge = Merge()
local_upstream_task = local_flow.get_tasks(name='trigger').pop()
qtzal_upstream_task = qtzal_flow.get_tasks(name='trigger').pop()
local_downstream_task = local_flow.reference_tasks().pop()
qtzal_downstream_task = qtzal_flow.reference_tasks().pop()
# Define flow and its task connections
with self:
data_source = Parameter('data_source',
default=data_source,
required=False)
switch(condition=data_source,
cases={
'local': local_upstream_task,
'quetzal': qtzal_upstream_task,
})
dataset = merge(local_branch=local_downstream_task,
quetzal_branch=qtzal_downstream_task)
self.set_reference_tasks([dataset])
logger.debug('Built flow %s with tasks %s', self, self.tasks)
def test_switch(condition_value):
condition = Condition()
a_branch = SuccessTask(name="a")
b_branch = SuccessTask(name="b")
c_branch = SuccessTask(name="c")
d_branch = SuccessTask(name="d")
with Flow(name="test") as flow:
switch(condition, dict(a=a_branch, b=b_branch, c=c_branch, d=d_branch))
assert len(flow.tasks) == 9
with prefect.context(CONDITION=condition_value):
state = flow.run()
assert isinstance(
state.result[a_branch], Success if condition_value == "a" else Skipped
)
assert isinstance(
state.result[b_branch], Success if condition_value == "b" else Skipped
)
assert isinstance(
state.result[c_branch], Success if condition_value == "c" else Skipped
)
assert isinstance(
state.result[d_branch], Success if condition_value == "d" else Skipped
)
def test_switch(condition_value):
condition = Condition()
with Flow(name="test") as flow:
a_branch = identity("a_val")
b_branch = identity("b_val")
c_branch = identity("c_val")
d_branch = identity("d_val")
res = switch(condition,
dict(a=a_branch, b=b_branch, c=c_branch, d=d_branch))
assert len(flow.tasks) == 10
with prefect.context(CONDITION=condition_value):
state = flow.run()
assert isinstance(state.result[a_branch],
Success if condition_value == "a" else Skipped)
assert isinstance(state.result[b_branch],
Success if condition_value == "b" else Skipped)
assert isinstance(state.result[c_branch],
Success if condition_value == "c" else Skipped)
assert isinstance(state.result[d_branch],
Success if condition_value == "d" else Skipped)
if condition_value == "x":
assert isinstance(state.result[res], Skipped)
else:
assert state.result[res].result == f"{condition_value}_val"
def test_mapped_switch_and_merge():
with Flow("iterated map") as flow:
mapped_result = identity.copy().map(["a", "b", "c"])
a = identity("a")
b = identity("b")
c = identity("c")
switch(condition=mapped_result, cases=dict(a=a, b=b, c=c), mapped=True)
merge_result = merge(a, b, c, mapped=True)
state = flow.run()
assert state.result[a].result == ["a", None, None]
assert state.result[b].result == [None, "b", None]
assert state.result[c].result == [None, None, "c"]
assert state.result[merge_result].result == ["a", "b", "c"]
def test_switch_works_with_raise_on_exception():
@prefect.task
def return_b():
return "b"
tasks = {let: prefect.Task(name=let) for let in "abcde"}
with Flow(name="test") as flow:
res = switch(return_b, tasks)
with raise_on_exception():
flow_state = flow.run()
logger = prefect.context.get("logger")
x, op, y = expression.split(' ')
logger.info("Received {} {} {}".format(x, op, y))
return dict(x=float(x), op=op, y=float(y))
# 'Arithmetic' is the name of the flow
with Flow('Arithmetic') as flow:
inputs = parse_input(Parameter('expression'))
# once we have our inputs, we create a dict of operations:
x, y = inputs['x'], inputs['y']
operations = {'+': x + y, '-': x - y, '*': x * y, '/': x / y}
# use prefect's `switch` task (conditional) to branch the flow, selecting the right operation
switch(condition=inputs['op'], cases=operations)
# use prefect's `merge` task to bring bring the branches from `switch` back together, producing a result
result = merge(*operations.values())
# do something with the result...
Result(result)
task_logger = get_logger("Task")
# Register the `Arithmetic` flow in the `Test1` project
flow.register(project_name='Test1')
# flow.run()
# print(flow.serialize())
@task
def c_task():
print("CCCCCC!")
@task
def check1():
result = False
print("check1", result)
return result
@task
def check2():
result = False
print("check2", result)
return result
@task(skip_on_upstream_skip=False)
def finish():
print("finish!")
with Flow("My Flow") as flow:
conditionally_run_a = switch(check1, {True: a_task})
conditionally_run_c = switch(check2, {True: b_task})
flow.chain(start, a_task, b_task, c_task, finish)
flow.run()
from prefect import task, Flow
from prefect.tasks.control_flow import switch
@task
def condition():
return "b"
@task
def a_branch():
return "A Branch"
@task
def b_branch():
return "B Branch"
with Flow("switch-flow") as flow:
switch(condition, dict(a=a_branch, b=b_branch))
flow.run()
with Flow("test") as flow:
producer = identity.copy().bind(["a", "b"])
cond = identity.copy().bind(producer, mapped=True)
a = identity.copy().bind("a")
a.set_upstream(producer, mapped=True)
b = identity.copy().bind("b")
b.set_upstream(producer, mapped=True)
c = identity.copy().bind("c")
c.set_upstream(producer, mapped=True)
switch(cond, cases=dict(a=a, b=b, c=c), mapped=True)
d = merge(a, b, mapped=True)
# state = flow.run()
# assert state.result[cond].result == ["a", "b"]
# assert state.result[a].result == ["a", None]
# assert state.result[b].result == [None, "b"]
# assert state.result[c].result == [None, None]
# assert state.result[d].result == ["a", "b"]
@task
def get_data():
return [1, 2, 3, 4]
