from metaflow import Parameter, current, Flow, step, FlowSpec
from functools import wraps

def analyze_artifacts(f):
    @wraps(f)
    def func(self):
        if f.__name__ == 'start':
            self._base_artifacts = {x for x in dir(self)}
            self._base_artifacts.add('_base_artifacts')
            inputs = set()
        else:
            inputs = {x for x in dir(self) if x not in self._base_artifacts}
        print('input artifacts', inputs)
        f(self)
        outputs = {x for x in dir(self) if x not in self._base_artifacts and x not in inputs}
        print('output artifacts', outputs)
    return func

class AnalyzeArtifactFlow(FlowSpec):

    param = Parameter('param')

    @analyze_artifacts
    @step
    def start(self):
        self.first = 'foo'
        self.second = 'bar'
        self.next(self.middle)

    @analyze_artifacts
    @step
    def middle(self):
        print('reading', self.first)
        x = self.param
        self.third = 'xyz'   
        self.next(self.end)

    @analyze_artifacts
    @step
    def end(self):
        pass

if __name__ == '__main__':
    AnalyzeArtifactFlow()