53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
try:
|
|
from UserDict import DictMixin
|
|
except ImportError:
|
|
from collections import MutableMapping as DictMixin
|
|
|
|
import os
|
|
|
|
class EnvironmentVarGuard(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
|