[3.12] gh-131277: allow EnvironmentVarGuard to unset more than one environment variable at once (GH-131280) (#131410)

picnixz · hugovk · web-flow · commit 865bd6de991d · 2025-03-18T13:46:24.000+01:00
(cherry picked from commit 3185e31) --------- Co-authored-by: Hugo van Kemenade <1324225+hugovk@users.noreply.github.com>
diff --git a/Lib/test/support/os_helper.py b/Lib/test/support/os_helper.py
@@ -676,9 +676,10 @@ def temp_umask(umask):
 
 
 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
@@ -712,8 +713,10 @@ def __len__(self):
     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.
diff --git a/Lib/test/test__osx_support.py b/Lib/test/test__osx_support.py
@@ -20,12 +20,13 @@ def setUp(self):
         self.prog_name = 'bogus_program_xxxx'
         self.temp_path_dir = os.path.abspath(os.getcwd())
         self.env = self.enterContext(os_helper.EnvironmentVarGuard())
-        for cv in ('CFLAGS', 'LDFLAGS', 'CPPFLAGS',
-                            'BASECFLAGS', 'BLDSHARED', 'LDSHARED', 'CC',
-                            'CXX', 'PY_CFLAGS', 'PY_LDFLAGS', 'PY_CPPFLAGS',
-                            'PY_CORE_CFLAGS', 'PY_CORE_LDFLAGS'):
-            if cv in self.env:
-                self.env.unset(cv)
+
+        self.env.unset(
+            'CFLAGS', 'LDFLAGS', 'CPPFLAGS',
+            'BASECFLAGS', 'BLDSHARED', 'LDSHARED', 'CC',
+            'CXX', 'PY_CFLAGS', 'PY_LDFLAGS', 'PY_CPPFLAGS',
+            'PY_CORE_CFLAGS', 'PY_CORE_LDFLAGS'
+        )
 
     def add_expected_saved_initial_values(self, config_vars, expected_vars):
         # Ensure that the initial values for all modified config vars
diff --git a/Lib/test/test_getopt.py b/Lib/test/test_getopt.py
@@ -13,8 +13,7 @@
 class GetoptTests(unittest.TestCase):
     def setUp(self):
         self.env = self.enterContext(EnvironmentVarGuard())
-        if "POSIXLY_CORRECT" in self.env:
-            del self.env["POSIXLY_CORRECT"]
+        del self.env["POSIXLY_CORRECT"]
 
     def assertError(self, *args, **kwargs):
         self.assertRaises(getopt.GetoptError, *args, **kwargs)
diff --git a/Lib/test/test_pathlib.py b/Lib/test/test_pathlib.py
@@ -3111,7 +3111,7 @@ def test_expanduser(self):
         p7 = P(f'~{fakename}/Documents')
 
         with os_helper.EnvironmentVarGuard() as env:
-            env.pop('HOME', None)
+            env.unset('HOME')
 
             self.assertEqual(p1.expanduser(), P(userhome) / 'Documents')
             self.assertEqual(p2.expanduser(), P(userhome) / 'Documents')
@@ -3222,10 +3222,7 @@ def test_rglob(self):
     def test_expanduser(self):
         P = self.cls
         with os_helper.EnvironmentVarGuard() as env:
-            env.pop('HOME', None)
-            env.pop('USERPROFILE', None)
-            env.pop('HOMEPATH', None)
-            env.pop('HOMEDRIVE', None)
+            env.unset('HOME', 'USERPROFILE', 'HOMEPATH', 'HOMEDRIVE')
             env['USERNAME'] = 'alice'
 
             # test that the path returns unchanged
@@ -3263,8 +3260,7 @@ def check():
             env['HOMEPATH'] = 'Users\\alice'
             check()
 
-            env.pop('HOMEDRIVE', None)
-            env.pop('HOMEPATH', None)
+            env.unset('HOMEDRIVE', 'HOMEPATH')
             env['USERPROFILE'] = 'C:\\Users\\alice'
             check()
 
diff --git a/Lib/test/test_platform.py b/Lib/test/test_platform.py
@@ -304,8 +304,7 @@ def raises_oserror(*a):
         with support.swap_attr(platform, '_wmi_query', raises_oserror):
             with os_helper.EnvironmentVarGuard() as environ:
                 try:
-                    if 'PROCESSOR_ARCHITEW6432' in environ:
-                        del environ['PROCESSOR_ARCHITEW6432']
+                    del environ['PROCESSOR_ARCHITEW6432']
                     environ['PROCESSOR_ARCHITECTURE'] = 'foo'
                     platform._uname_cache = None
                     system, node, release, version, machine, processor = platform.uname()
diff --git a/Lib/test/test_regrtest.py b/Lib/test/test_regrtest.py
@@ -409,8 +409,7 @@ def create_regrtest(self, args):
         # which has an unclear API
         with os_helper.EnvironmentVarGuard() as env:
             # Ignore SOURCE_DATE_EPOCH env var if it's set
-            if 'SOURCE_DATE_EPOCH' in env:
-                del env['SOURCE_DATE_EPOCH']
+            del env['SOURCE_DATE_EPOCH']
 
             regrtest = main.Regrtest(ns)
 
diff --git a/Lib/test/test_shutil.py b/Lib/test/test_shutil.py
@@ -2377,7 +2377,7 @@ def test_environ_path_cwd(self):
 
     def test_environ_path_missing(self):
         with os_helper.EnvironmentVarGuard() as env:
-            env.pop('PATH', None)
+            del env['PATH']
 
             # without confstr
             with unittest.mock.patch('os.confstr', side_effect=ValueError, \
@@ -2403,7 +2403,7 @@ def test_empty_path(self):
 
     def test_empty_path_no_PATH(self):
         with os_helper.EnvironmentVarGuard() as env:
-            env.pop('PATH', None)
+            del env['PATH']
             rv = shutil.which(self.file)
             self.assertIsNone(rv)
 
@@ -3335,17 +3335,15 @@ def test_stty_match(self):
         expected = (int(size[1]), int(size[0])) # reversed order
 
         with os_helper.EnvironmentVarGuard() as env:
-            del env['LINES']
-            del env['COLUMNS']
+            env.unset('LINES', 'COLUMNS')
             actual = shutil.get_terminal_size()
 
         self.assertEqual(expected, actual)
 
     @unittest.skipIf(support.is_wasi, "WASI has no /dev/null")
     def test_fallback(self):
         with os_helper.EnvironmentVarGuard() as env:
-            del env['LINES']
-            del env['COLUMNS']
+            env.unset('LINES', 'COLUMNS')
 
             # sys.__stdout__ has no fileno()
             with support.swap_attr(sys, '__stdout__', None):
diff --git a/Lib/test/test_site.py b/Lib/test/test_site.py
@@ -354,9 +354,7 @@ def test_no_home_directory(self):
 
         with EnvironmentVarGuard() as environ, \
              mock.patch('os.path.expanduser', lambda path: path):
-
-            del environ['PYTHONUSERBASE']
-            del environ['APPDATA']
+            environ.unset('PYTHONUSERBASE', 'APPDATA')
 
             user_base = site.getuserbase()
             self.assertTrue(user_base.startswith('~' + os.sep),
