diff --git a/Lib/test/support/os_helper.py b/Lib/test/support/os_helper.py index d82093e375c8b1..5e2d673b544f1f 100644 --- a/Lib/test/support/os_helper.py +++ b/Lib/test/support/os_helper.py @@ -7,6 +7,7 @@ import stat import string import sys +import threading import time import unittest import warnings @@ -736,64 +737,73 @@ def temp_umask(umask): class EnvironmentVarGuard(collections.abc.MutableMapping): - """Class to help protect the environment variable properly. - + """Thread-safe class to help protect environment variables. Can be used as a context manager. """ def __init__(self): self._environ = os.environ self._changed = {} + self._lock = threading.RLock() def __getitem__(self, envvar): - return self._environ[envvar] + with self._lock: + 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 + with self._lock: + # 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] + with self._lock: + # Remember the initial value on the first access + if envvar not in self._changed: + self._changed[envvar] = self._environ.get(envvar) + self._environ.pop(envvar, None) def keys(self): - return self._environ.keys() + with self._lock: + return list(self._environ.keys()) def __iter__(self): - return iter(self._environ) + with self._lock: + return iter(dict(self._environ)) def __len__(self): - return len(self._environ) + with self._lock: + return len(self._environ) def set(self, envvar, value): self[envvar] = value def unset(self, envvar, /, *envvars): """Unset one or more environment variables.""" - for ev in (envvar, *envvars): - del self[ev] + with self._lock: + for ev in (envvar, *envvars): + del self[ev] def copy(self): - # We do what os.environ.copy() does. - return dict(self) + with self._lock: + return dict(self._environ) 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 + with self._lock: + for (k, v) in self._changed.items(): + if v is None: + self._environ.pop(k, None) + else: + self._environ[k] = v + self._changed.clear() + os.environ = self._environ + def __reduce__(self): + return (dict, (dict(self),)) try: if support.MS_WINDOWS: diff --git a/Lib/test/test_support.py b/Lib/test/test_support.py index efe6b77a7faa18..5c8214b3e4a9c8 100644 --- a/Lib/test/test_support.py +++ b/Lib/test/test_support.py @@ -13,6 +13,8 @@ import sysconfig import tempfile import textwrap +import threading +import time import unittest import warnings @@ -22,6 +24,7 @@ from test.support import script_helper from test.support import socket_helper from test.support import warnings_helper +from test.support.os_helper import EnvironmentVarGuard TESTFN = os_helper.TESTFN @@ -794,6 +797,27 @@ def test_linked_to_musl(self): for v in linked: self.assertIsInstance(v, int) + def test_threadsafe_environmentvarguard(self): + def worker1(guard): + for i in range(1000): + guard['MY_VAR'] = 'value1' + time.sleep(0.0001) # Small delay to increase chance of thread switching + + def worker2(guard): + for i in range(1000): + guard['MY_VAR'] = 'value2' + time.sleep(0.0001) + + with EnvironmentVarGuard() as guard: + t1 = threading.Thread(target=worker1, args=(guard,)) + t2 = threading.Thread(target=worker2, args=(guard,)) + t1.start() + t2.start() + t1.join() + t2.join() + final_value = os.getenv('MY_VAR') + self.assertIn(final_value, ("value1", "value2")) + # XXX -follows a list of untested API # make_legacy_pyc diff --git a/Misc/NEWS.d/next/Tests/2025-04-07-19-30-16.gh-issue-132113.GIgcmV.rst b/Misc/NEWS.d/next/Tests/2025-04-07-19-30-16.gh-issue-132113.GIgcmV.rst new file mode 100644 index 00000000000000..5ae72558841932 --- /dev/null +++ b/Misc/NEWS.d/next/Tests/2025-04-07-19-30-16.gh-issue-132113.GIgcmV.rst @@ -0,0 +1 @@ +Make :class:`EnvironmentVarGuard ` thread safe