from time import sleep
from metaflow import step, FlowSpec, catch, timeout 

TIMEOUT=60

class CatchTimeoutForeachFlow(FlowSpec):
    @step
    def start(self):
        print("Starting 👋")
        self.items = [0,1]
        self.next(self.work, foreach="items")
    
    @catch(var="timeout_error")
    @timeout(seconds=TIMEOUT)
    @step
    def work(self):
        print("Starting 👋")
        sleep(TIMEOUT+10)
        self.next(self.join)

    @step
    def join(self, inputs):
        aggr = []
        for input in inputs:
            if not getattr(input, "timeout_error"):
                continue
            aggr.append(input.timeout_error)
        
        if aggr:
            self.timeout_error = aggr[0]
                
        self.next(self.end)

    @step
    def end(self):
        print("Done! 🏁")

        print(f"timeout_error is : {self.timeout_error}")


if __name__ == "__main__":
    CatchTimeoutForeachFlow()