Import a function from one python file and use that function as a step function in metaflow

I have a function sayHello defined in one python file utils.py. I also have a python file flow.py that runs a simple metaflow. I want to import sayHello from utils.py and use sayHello as a step function in flow.py. Is it possible to do that? This might be hard because we need self in the class. If it’s possible, one question is that we how do pass output from the previous step into the function and pass the output to next step. The following is my attempt.

#utils.py
def sayHello():
    print("hello world")

#flow.py
from metaflow import FlowSpec, step, Parameter
from metaflow import Metaflow
from utils import sayHello

def function(p):
    return p

class BranchFlow(FlowSpec):
    @step
    def start(self):
        self.next(self.a, self.b, self.sayHello)

    @step
    def a(self):
        self.x = 1
        self.next(self.join)

    @step
    def b(self):
        self.x = 2
        self.next(self.join)


    @step
    def join(self, inputs):
        print('a is %s' % inputs.a.x)
        print('b is %s' % inputs.b.x)
        print('total is %d' % sum(input.x for input in inputs))
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == '__main__':
     BranchFlow()    

>Solution :

According to the documentation for FlowSpec.next(), the method you pass to self.next() must be a member of the current class and decorated with @step.

Probably the easiest way to use another function from another module is to wrap it in a member function:

class BranchFlow(FlowSpec):
    # ...

    @step
    def sayHello(self):
        sayHello()

Now you can do self.next(self.sayHello).