A simple set of classes for dealing with environment variables.

EnvSetting.py

__author__ = 'Steve'

import os


class EnvSetting(object):
    """
    Set and unset an environemnt variable. If the variable did not exist at set time, it will be deleted at unset time

    setter = EnvSetting('var', value)
    setter.set()
    do_something_with_var()
    setter.unset()
    """
    def __init__(self, key, val):
        self._key = key
        self._val = str(val)
        self._cache = None

    def _restore(self, val):
        os.environ[self._key] = val

    def _delete(self):
        del os.environ[self._key]

    def set(self):
        if self._key in os.environ:
            current = os.environ.get(self._key)
            self._cache = lambda: self._restore(current)
        else:
            self._cache = lambda: self._delete()

        os.environ[self._key] = self._val

    def unset(self):
        self._cache()

    def __repr__(self):
        return "(os.environ['%s'] = '%s')" % (self._key, self._val)


class EnvCtx(object):
    """
    A context manager that sets an environment variable and unsets it at exit time

        with EnvCtx('var', value):
             do_something_with_var()

    """
    def __init__(self, name, val):
        self.setter = EnvSetting(name, val)

    def __enter__(self):
        self.setter.set()

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.setter.unset()


class EnvSetDecorator(object):
    """
    A decorated that applies an EnvCtx context to the decorated function

        @EnvSetDecorator('var', value)
        def function_that_needs_var():
           ....

    """
    def __init__(self, name, val):
        self.context = EnvCtx(name, val)

    def __call__(self, fn):
        def wrapped(*args, **kwargs):
            with self.context:
                return fn(*args, **kwargs)

        return wrapped