diff --git a/src/python/pants/util/contextutil.py b/src/python/pants/util/contextutil.py index bf6c3baeada..3f64e9d6798 100644 --- a/src/python/pants/util/contextutil.py +++ b/src/python/pants/util/contextutil.py @@ -48,6 +48,16 @@ def get_joined_path(new_entries, env=None, env_var='PATH', delimiter=':', prepen return delimiter.join(path_dirs) +def _os_encode(u, enc=sys.getfilesystemencoding()): + """Turns a `unicode` into `bytes` via encoding.""" + return u.encode(enc, 'strict') + + +def _os_decode(b, enc=sys.getfilesystemencoding()): + """Turns a `bytes` into `unicode` via decoding.""" + return b.decode(enc, 'strict') + + @contextmanager def environment_as(**kwargs): """Update the environment to the supplied values, for example: @@ -61,7 +71,7 @@ def environment_as(**kwargs): def setenv(key, val): if val is not None: - os.environ[key] = val + os.environ[key] = _os_encode(val) else: if key in os.environ: del os.environ[key] @@ -76,6 +86,10 @@ def setenv(key, val): setenv(key, val) +def _copy_and_decode_env(env): + return {k: _os_decode(v) for k, v in env.items()} + + 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 @@ -87,13 +101,13 @@ def _purge_env(): def _restore_env(env): for k, v in env.items(): - os.environ[k] = v + os.environ[k] = _os_encode(v) @contextmanager def hermetic_environment_as(**kwargs): """Set the environment to the supplied values from an empty state.""" - old_environment = os.environ.copy() + old_environment = _copy_and_decode_env(os.environ) _purge_env() try: with environment_as(**kwargs): diff --git a/tests/python/pants_test/pantsd/test_pantsd_integration.py b/tests/python/pants_test/pantsd/test_pantsd_integration.py index 99e18f0c855..ed1f1763639 100644 --- a/tests/python/pants_test/pantsd/test_pantsd_integration.py +++ b/tests/python/pants_test/pantsd/test_pantsd_integration.py @@ -549,3 +549,11 @@ def test_pantsd_environment_scrubbing(self): result = pantsd_run(['compile', 'examples/src/java/org/pantsbuild/example/hello/simple']) self.assert_success(result) + + def test_pantsd_unicode_environment(self): + with self.pantsd_successful_run_context( + extra_env={'XXX': '¡'}, + ) as (pantsd_run, checker, workdir, _): + result = pantsd_run(['help']) + checker.assert_started() + self.assert_success(result) diff --git a/tests/python/pants_test/util/test_contextutil.py b/tests/python/pants_test/util/test_contextutil.py index 03b1cda71b9..6a49835b1fc 100644 --- a/tests/python/pants_test/util/test_contextutil.py +++ b/tests/python/pants_test/util/test_contextutil.py @@ -74,6 +74,16 @@ def test_hermetic_environment_subprocesses(self): self.assertIn('USER', os.environ) self.assertNotIn('AAA', os.environ) + def test_hermetic_environment_unicode(self): + UNICODE_CHAR = '¡' + ENCODED_CHAR = UNICODE_CHAR.encode('utf-8') + with environment_as(**dict(XXX=UNICODE_CHAR)): + self.assertEquals(os.environ['XXX'], ENCODED_CHAR) + with hermetic_environment_as(**dict(AAA=UNICODE_CHAR)): + self.assertIn('AAA', os.environ) + self.assertEquals(os.environ['AAA'], ENCODED_CHAR) + self.assertEquals(os.environ['XXX'], ENCODED_CHAR) + def test_simple_pushd(self): pre_cwd = os.getcwd() with temporary_dir() as tempdir: