49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
|
import UserDict
|
||
|
|
import os
|
||
|
|
|
||
|
|
class EnvironmentVarGuard(UserDict.DictMixin):
|
||
|
|
|
||
|
|
"""Class to help protect the environment variable properly. Can be used as
|
||
|
|
a context manager."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._environ = os.environ
|
||
|
|
self._changed = {}
|
||
|
|
|
||
|
|
def __getitem__(self, envvar):
|
||
|
|
return self._environ[envvar]
|
||
|
|
|
||
|
|
def __setitem__(self, envvar, value):
|
||
|
|
# Remember the initial value on the first access
|
||
|
|
if envvar not in self._changed:
|
||
|
|
self._changed[envvar] = self._environ.get(envvar)
|
||
|
|
self._environ[envvar] = value
|
||
|
|
|
||
|
|
def __delitem__(self, envvar):
|
||
|
|
# Remember the initial value on the first access
|
||
|
|
if envvar not in self._changed:
|
||
|
|
self._changed[envvar] = self._environ.get(envvar)
|
||
|
|
if envvar in self._environ:
|
||
|
|
del self._environ[envvar]
|
||
|
|
|
||
|
|
def keys(self):
|
||
|
|
return self._environ.keys()
|
||
|
|
|
||
|
|
def set(self, envvar, value):
|
||
|
|
self[envvar] = value
|
||
|
|
|
||
|
|
def unset(self, envvar):
|
||
|
|
del self[envvar]
|
||
|
|
|
||
|
|
def __enter__(self):
|
||
|
|
return self
|
||
|
|
|
||
|
|
def __exit__(self, *ignore_exc):
|
||
|
|
for (k, v) in self._changed.items():
|
||
|
|
if v is None:
|
||
|
|
if k in self._environ:
|
||
|
|
del self._environ[k]
|
||
|
|
else:
|
||
|
|
self._environ[k] = v
|
||
|
|
os.environ = self._environ
|