from typing import List

from prefect import Flow, task


@task(checkpoint=False)
def task_1() -> List[int]:
    return list(range(15))


@task(checkpoint=False)
def task_2(input: int) -> int:
    return input + 1


@task(checkpoint=False)
def task_3(list_of_ints: List[int]) -> str:
    print(f"---------------- len: {len(list_of_ints)} -----------------")
    return "task_3 done"


@task(checkpoint=False)
def task_4(input: int) -> int:
    return input + 2


with Flow("Flow") as flow:
    list_of_ints_1 = task_1()
    list_of_ints_2 = task_2.map(list_of_ints_1)
    str_output = task_3(list_of_ints_2)
    task_4.map(list_of_ints_2, upstream_tasks=[str_output])

status = flow.run()
flow.visualize(status, filename="./flow", format="svg")