diff options
Diffstat (limited to 'Lib/test/support')
| -rw-r--r-- | Lib/test/support/__init__.py | 3 | ||||
| -rw-r--r-- | Lib/test/support/os_helper.py | 11 |
2 files changed, 8 insertions, 6 deletions
diff --git a/Lib/test/support/__init__.py b/Lib/test/support/__init__.py index b9ccf7bb4c6..b35ce1b2285 100644 --- a/Lib/test/support/__init__.py +++ b/Lib/test/support/__init__.py @@ -2855,8 +2855,7 @@ def no_color(): swap_attr(_colorize, "can_colorize", lambda file=None: False), EnvironmentVarGuard() as env, ): - for var in {"FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS"}: - env.unset(var) + env.unset("FORCE_COLOR", "NO_COLOR", "PYTHON_COLORS") env.set("NO_COLOR", "1") yield diff --git a/Lib/test/support/os_helper.py b/Lib/test/support/os_helper.py index 15dcdc9b1fd..4e39b5a835a 100644 --- a/Lib/test/support/os_helper.py +++ b/Lib/test/support/os_helper.py @@ -720,9 +720,10 @@ else: class EnvironmentVarGuard(collections.abc.MutableMapping): + """Class to help protect the environment variable properly. - """Class to help protect the environment variable properly. Can be used as - a context manager.""" + Can be used as a context manager. + """ def __init__(self): self._environ = os.environ @@ -756,8 +757,10 @@ class EnvironmentVarGuard(collections.abc.MutableMapping): def set(self, envvar, value): self[envvar] = value - def unset(self, envvar): - del self[envvar] + def unset(self, envvar, /, *envvars): + """Unset one or more environment variables.""" + for ev in (envvar, *envvars): + del self[ev] def copy(self): # We do what os.environ.copy() does. |