from metaflow import FlowSpec, step
from functools import wraps
import os

class strict_environment():

    def __init__(self, env=None):
        self.user = os.environ.get('USER')
        self.env = {} if env is None else env

    def __call__(self, f):
        @wraps(f)
        def func(s):
            s.original_environment = dict(os.environ)
            os.environ = self.env
            os.environ['USER'] = self.user
            f(s)
        return func

class StrictEnvironmentFlow(FlowSpec):

    @strict_environment(env={'SOME_ENVVAR': 'value'})
    @step
    def start(self):
        print("SOME_ENVVAR =", os.environ['SOME_ENVVAR'])
        self.next(self.end)

    @strict_environment(env={})
    @step
    def end(self):
        print('environment', os.environ)

if __name__ == '__main__':
    StrictEnvironmentFlow()