Tom Shaffner
01/21/2022, 12:04 AMKevin Kho
from prefect import Flow, task, case, apply_map
from prefect.tasks.control_flow import merge
@task
def inc(x):
return x + 1
@task
def negate(x):
return -x
@task
def is_even(x):
return x % 2 == 0
@task
def something():
return "end"
def inc_or_negate(x):
cond = is_even(x)
# If x is even, increment it
with case(cond, True):
res1 = inc(x)
# If x is odd, negate it
with case(cond, False):
res2 = negate(x)
return merge(res1, res2)
with Flow("apply-map example") as flow:
result = apply_map(inc_or_negate, range(4))
x = something()
x.set_upstream(result)
and schematic generated from a flow runTom Shaffner
01/21/2022, 3:27 PM