from prefect import flow, task
from prefect.states import Completed, Failed, State
from prefect.utilities.collections import listrepr

@task
def add_42(v: int) -> int:
    try:
        return v + 42
    except Exception as e:
        return Failed(message=f"eek!: {e!r}")

def handler(state) -> State:
    print(f"Handling state: {state!r}")
    return state
    
@task
def write_to_storage(v):
    print(f"Writing {v!r} to storage")

@flow(log_prints=True)
def foo():
    result_states = add_42.map([1, None, 3], return_state=True)

    completed, failed = (
        filter(lambda s: s.is_completed(), result_states),
        filter(lambda s: s.is_failed(), result_states),
    )
    
    failures = list(map(handler, failed))
    
    write_to_storage.map(completed)
    
    return Completed(message=f"Finished with some failures: {listrepr(failures)}")
    
if __name__ == "__main__":
    foo()