from metaflow import FlowSpec, step
from functools import wraps

class FakeSelf():

    def __init__(self, orig_self, key):
        self.orig_self = orig_self
        self.key = key

    def dataset_write(self):
        print('key', self.key)

    def __getattr__(self, k):
        return getattr(self.orig_self, k)

class extra_arg():

    def __init__(self, key="value"):
        self.key = key

    def __call__(self, f):
        @wraps(f)
        def func(s):
    	    return f(FakeSelf(s, self.key))
        return func

class FirstFlow(FlowSpec):

    @extra_arg(key="test")
    @step
    def start(self):
        self.dataset_write()
        self.next(self.end)

    @step
    def end(self):
        pass

if __name__ == '__main__':
    FirstFlow()