Issue #8603: Create a bytes version of os.environ for Unix

Create os.environb mapping and os.getenvb() function, os.unsetenv() encodes str
argument to the file system encoding with the surrogateescape error handler
(instead of utf8/strict) and accepts bytes, and posix.environ keys and values
are bytes.

3 files changed, 92 insertions, 24 deletions

diff --git a/Lib/os.py b/Lib/os.py
index 7672d6f..3e2ee0d 100644
--- a/Lib/os.py
+++ b/Lib/os.py
@@ -387,29 +387,33 @@ def get_exec_path(env=None):
 from _abcoll import MutableMapping  # Can't use collections (bootstrap)
 
 class _Environ(MutableMapping):
-    def __init__(self, environ, keymap, putenv, unsetenv):
-        self.keymap = keymap
+    def __init__(self, data, encodekey, decodekey, encodevalue, decodevalue, putenv, unsetenv):
+        self.encodekey = encodekey
+        self.decodekey = decodekey
+        self.encodevalue = encodevalue
+        self.decodevalue = decodevalue
         self.putenv = putenv
         self.unsetenv = unsetenv
-        self.data = data = {}
-        for key, value in environ.items():
-            data[keymap(key)] = str(value)
+        self.data = data
 
     def __getitem__(self, key):
-        return self.data[self.keymap(key)]
+        value = self.data[self.encodekey(key)]
+        return self.decodevalue(value)
 
     def __setitem__(self, key, value):
-        value = str(value)
+        key = self.encodekey(key)
+        value = self.encodevalue(value)
         self.putenv(key, value)
-        self.data[self.keymap(key)] = value
+        self.data[key] = value
 
     def __delitem__(self, key):
+        key = self.encodekey(key)
         self.unsetenv(key)
-        del self.data[self.keymap(key)]
+        del self.data[key]
 
     def __iter__(self):
         for key in self.data:
-            yield key
+            yield self.decodekey(key)
 
     def __len__(self):
         return len(self.data)
@@ -439,22 +443,67 @@ except NameError:
 else:
     __all__.append("unsetenv")
 
-if name in ('os2', 'nt'): # Where Env Var Names Must Be UPPERCASE
-    _keymap = lambda key: str(key.upper())
-else:  # Where Env Var Names Can Be Mixed Case
-    _keymap = lambda key: str(key)
-
-environ = _Environ(environ, _keymap, _putenv, _unsetenv)
+def _createenviron():
+    if name in ('os2', 'nt'):
+        # Where Env Var Names Must Be UPPERCASE
+        def check_str(value):
+            if not isinstance(value, str):
+                raise TypeError("str expected, not %s" % type(value).__name__)
+            return value
+        encode = check_str
+        decode = str
+        def encodekey(key):
+            return encode(key).upper()
+        data = {}
+        for key, value in environ.items():
+            data[encodekey(key)] = value
+    else:
+        # Where Env Var Names Can Be Mixed Case
+        def encode(value):
+            if not isinstance(value, str):
+                raise TypeError("str expected, not %s" % type(value).__name__)
+            return value.encode(sys.getfilesystemencoding(), 'surrogateescape')
+        def decode(value):
+            return value.decode(sys.getfilesystemencoding(), 'surrogateescape')
+        encodekey = encode
+        data = environ
+    return _Environ(data,
+        encodekey, decode,
+        encode, decode,
+        _putenv, _unsetenv)
+
+# unicode environ
+environ = _createenviron()
+del _createenviron
 
 
 def getenv(key, default=None):
     """Get an environment variable, return None if it doesn't exist.
-    The optional second argument can specify an alternate default."""
-    if isinstance(key, bytes):
-        key = key.decode(sys.getfilesystemencoding(), "surrogateescape")
+    The optional second argument can specify an alternate default.
+    key, default and the result are str."""
     return environ.get(key, default)
 __all__.append("getenv")
 
+if name not in ('os2', 'nt'):
+    def _check_bytes(value):
+        if not isinstance(value, bytes):
+            raise TypeError("bytes expected, not %s" % type(value).__name__)
+        return value
+
+    # bytes environ
+    environb = _Environ(environ.data,
+        _check_bytes, bytes,
+        _check_bytes, bytes,
+        _putenv, _unsetenv)
+    del _check_bytes
+
+    def getenvb(key, default=None):
+        """Get an environment variable, return None if it doesn't exist.
+        The optional second argument can specify an alternate default.
+        key, default and the result are bytes."""
+        return environb.get(key, default)
+    __all__.append("getenvb")
+
 def _exists(name):
     return name in globals()
 
diff --git a/Lib/test/test_os.py b/Lib/test/test_os.py
index b91f97b..49c6591 100644
--- a/Lib/test/test_os.py
+++ b/Lib/test/test_os.py
@@ -369,12 +369,15 @@ class EnvironTests(mapping_tests.BasicTestMappingProtocol):
 
     def setUp(self):
         self.__save = dict(os.environ)
+        self.__saveb = dict(os.environb)
         for key, value in self._reference().items():
             os.environ[key] = value
 
     def tearDown(self):
         os.environ.clear()
         os.environ.update(self.__save)
+        os.environb.clear()
+        os.environb.update(self.__saveb)
 
     def _reference(self):
         return {"KEY1":"VALUE1", "KEY2":"VALUE2", "KEY3":"VALUE3"}
@@ -439,6 +442,24 @@ class EnvironTests(mapping_tests.BasicTestMappingProtocol):
         # Supplied PATH environment variable
         self.assertSequenceEqual(test_path, os.get_exec_path(test_env))
 
+    @unittest.skipIf(sys.platform == "win32", "POSIX specific test")
+    def test_environb(self):
+        # os.environ -> os.environb
+        value = 'euro\u20ac'
+        try:
+            value_bytes = value.encode(sys.getfilesystemencoding(), 'surrogateescape')
+        except UnicodeEncodeError:
+            raise unittest.SkipTest("U+20AC character is not encodable to %s" % sys.getfilesystemencoding())
+        os.environ['unicode'] = value
+        self.assertEquals(os.environ['unicode'], value)
+        self.assertEquals(os.environb[b'unicode'], value_bytes)
+
+        # os.environb -> os.environ
+        value = b'\xff'
+        os.environb[b'bytes'] = value
+        self.assertEquals(os.environb[b'bytes'], value)
+        value_str = value.decode(sys.getfilesystemencoding(), 'surrogateescape')
+        self.assertEquals(os.environ['bytes'], value_str)
 
 class WalkTests(unittest.TestCase):
     """Tests for os.walk()."""
diff --git a/Lib/test/test_subprocess.py b/Lib/test/test_subprocess.py
index be163fc..eb96706 100644
--- a/Lib/test/test_subprocess.py
+++ b/Lib/test/test_subprocess.py
@@ -803,8 +803,6 @@ class POSIXProcessTestCase(BaseTestCase):
 
     def test_undecodable_env(self):
         for key, value in (('test', 'abc\uDCFF'), ('test\uDCFF', '42')):
-            value_repr = repr(value).encode("ascii")
-
             # test str with surrogates
             script = "import os; print(repr(os.getenv(%s)))" % repr(key)
             env = os.environ.copy()
@@ -813,19 +811,19 @@ class POSIXProcessTestCase(BaseTestCase):
                 [sys.executable, "-c", script],
                 env=env)
             stdout = stdout.rstrip(b'\n\r')
-            self.assertEquals(stdout, value_repr)
+            self.assertEquals(stdout.decode('ascii'), repr(value))
 
             # test bytes
             key = key.encode("ascii", "surrogateescape")
             value = value.encode("ascii", "surrogateescape")
-            script = "import os; print(repr(os.getenv(%s)))" % repr(key)
+            script = "import os; print(repr(os.getenvb(%s)))" % repr(key)
             env = os.environ.copy()
             env[key] = value
             stdout = subprocess.check_output(
                 [sys.executable, "-c", script],
                 env=env)
             stdout = stdout.rstrip(b'\n\r')
-            self.assertEquals(stdout, value_repr)
+            self.assertEquals(stdout.decode('ascii'), repr(value))
 
 
 @unittest.skipUnless(mswindows, "Windows specific tests")