diff --git a/src/python/pants/util/contextutil.py b/src/python/pants/util/contextutil.py index d904ea2c8a4..bf6c3baeada 100644 --- a/src/python/pants/util/contextutil.py +++ b/src/python/pants/util/contextutil.py @@ -76,15 +76,31 @@ def setenv(key, val): setenv(key, val) +def _purge_env(): + # N.B. Without the use of `del` here (which calls `os.unsetenv` under the hood), subprocess32 + # invokes or other things that may access the environment at the C level may not see the + # correct env vars (i.e. we can't just replace os.environ with an empty dict). + # See https://docs.python.org/2/library/os.html#os.unsetenv for more info. + for k in os.environ.keys(): + del os.environ[k] + + +def _restore_env(env): + for k, v in env.items(): + os.environ[k] = v + + @contextmanager def hermetic_environment_as(**kwargs): """Set the environment to the supplied values from an empty state.""" - old_environment, os.environ = os.environ, {} + old_environment = os.environ.copy() + _purge_env() try: with environment_as(**kwargs): yield finally: - os.environ = old_environment + _purge_env() + _restore_env(old_environment) @contextmanager diff --git a/tests/python/pants_test/util/test_contextutil.py b/tests/python/pants_test/util/test_contextutil.py index 2a018636647..03b1cda71b9 100644 --- a/tests/python/pants_test/util/test_contextutil.py +++ b/tests/python/pants_test/util/test_contextutil.py @@ -64,6 +64,16 @@ def test_hermetic_environment(self): with hermetic_environment_as(**{}): self.assertNotIn('USER', os.environ) + def test_hermetic_environment_subprocesses(self): + self.assertIn('USER', os.environ) + with hermetic_environment_as(**dict(AAA='333')): + output = subprocess.check_output('env', shell=True) + self.assertNotIn('USER=', output) + self.assertIn('AAA', os.environ) + self.assertEquals(os.environ['AAA'], '333') + self.assertIn('USER', os.environ) + self.assertNotIn('AAA', os.environ) + def test_simple_pushd(self): pre_cwd = os.getcwd() with temporary_dir() as tempdir: