diff --git a/pydu/environ.py b/pydu/environ.py new file mode 100644 index 0000000..5d4ec73 --- /dev/null +++ b/pydu/environ.py @@ -0,0 +1,27 @@ +import os +from pydu.compat import iteritems +from contextlib import contextmanager + + +@contextmanager +def environ(**kwargs): + """ + Context manager for updating one or more environment variables. + + Preserves the previous environment variable (if available) and + recovers when exiting the context manager. + """ + original_kwargs = {} + + for key in kwargs: + original_kwargs[key] = os.environ.get(key, None) + os.environ[key] = kwargs[key] + + yield + + for key, value in iteritems(original_kwargs): + if value is None: + del os.environ[key] + continue + + os.environ[key] = value diff --git a/tests/test_environ.py b/tests/test_environ.py new file mode 100644 index 0000000..7ec3ebf --- /dev/null +++ b/tests/test_environ.py @@ -0,0 +1,10 @@ +import os +from pydu.environ import environ + + +def test_environ(): + with environ(a='a', b=''): + assert os.environ['a'] == 'a' + assert os.environ['b'] == '' + assert 'a' not in os.environ + assert 'b' not in os.environ