Add environ and environb

CHANGES

@@ -12,6 +12,10 @@ Development version
 - Pull request #204: Add `six.ensure_binary`, `six.ensure_text`, and
   `six.ensure_str`.
 
+- Add `six.supports_bytes_environ`, `six.environ`, and `six.environb` (when
+  `six.supports_bytes_environ` is True) which operate similar to the
+  Python-3.2+ `os.supports_bytes_environ`,  `os.environ`, and `os.environb`.
+
 1.11.0
 ------
 

six.py

@@ -25,6 +25,7 @@
 import functools
 import itertools
 import operator
+import os
 import sys
 import types
 
@@ -621,6 +622,88 @@ def iterlists(d, **kw):
          "Return an iterator over the (key, [values]) pairs of a dictionary.")
 
 
+if PY3:
+    environ = os.environ
+    if sys.version_info >= (3, 2):
+        # We can simply reuse what Python 3.2+ provides
+        supports_bytes_environ = os.supports_bytes_environ
+        if supports_bytes_environ:
+            environb = os.environb
+    else:
+        # There is no os.environb on Python 3.0 and 3.1
+        supports_bytes_environ = False
+else:
+    # Python 2's environ is bytes-oriented
+    supports_bytes_environ = True
+    environb = os.environ
+    if sys.version_info >= (2, 6):
+        # Python 2.6 and Python 2.7 have MutableMapping
+        from collections import MutableMapping as DictBaseClass
+    else:
+        # Python 2.5 and below do not have MutableMapping
+        from UserDict import DictMixin as DictBaseClass
+
+    class _TextEnviron(DictBaseClass):
+        """
+        Utility class to return text strings from the environment instead of byte strings
+
+        Mimics the behaviour of os.environ on Python3
+        """
+        # Mimic Python3's os.environ by using sys.getfilesystemencoding()
+        def __init__(self, env=None, encoding=sys.getfilesystemencoding()):
+            if env is None:
+                env = os.environ
+            self._raw_environ = env
+            self._value_cache = {}
+            self.encoding = encoding
+
+        def __delitem__(self, key):
+            del self._raw_environ[key]
+
+        def __getitem__(self, key):
+            # Note: For undecodable strings, Python3 will use surrogateescape.  We don't have that
+            # on Python2 so we throw a ValueError instead
+            value = self._raw_environ[key]
+
+            # Cache keys off of the undecoded values to handle any environment variables which
+            # change during a run
+            if value not in self._value_cache:
+                try:
+                    self._value_cache[value] = ensure_text(value, encoding=self.encoding)
+                except UnicodeError:
+                    raise ValueError('environ string for %s is undecodable in the'
+                                     ' filesystemencoding. Use six.environb and manually convert'
+                                     ' the value to text instead' % value)
+            return self._value_cache[value]
+
+        def __setitem__(self, key, value):
+            if not isinstance(value, text_type):
+                raise TypeError("str expected, not %s" % type(value).__name__)
+
+            try:
+                self._raw_environ[key] = ensure_binary(value, encoding=self.encoding)
+            except UnicodeError:
+                raise ValueError('The value, %s, is unencodable in the filesystemencoding.'
+                                 ' Use six.environb and manually convert the value to bytes'
+                                 ' instead' % value)
+
+        def __iter__(self):
+            return self._raw_environ.__iter__()
+
+        def __len__(self):
+            return len(self._raw_environ)
+
+        # Needed for the DictMixin-based subclass but not for MutableMapping-based subclass
+        if sys.version_info < (2, 6):
+            def keys(self):
+                return self._raw_environ.keys()
+
+            def copy(self):
+                return self._raw_environ.copy()
+
+    environ = _TextEnviron()
+
+
 if PY3:
     def b(s):
         return s.encode("latin-1")
@@ -908,7 +991,6 @@ def ensure_text(s, encoding='utf-8', errors='strict'):
         raise TypeError("not expecting type '%s'" % type(s))
 
 
-
 def python_2_unicode_compatible(klass):
     """
     A decorator that defines __unicode__ and __str__ methods under Python 2.

test_six.py

@@ -1010,3 +1010,57 @@ def test_ensure_text(self):
             assert converted_unicode == self.UNICODE_EMOJI and isinstance(converted_unicode, str)
             # PY3: bytes -> str
             assert converted_binary == self.UNICODE_EMOJI and isinstance(converted_unicode, str)
+
+
+@py.test.fixture
+def env_var():
+    try:
+        old_six_test = os.environ['_six_test']
+    except KeyError:
+        old_six_test = None
+    else:
+        del os.environ['_six_test']
+
+    yield
+
+    if old_six_test is not None:
+        os.environ['_six_test'] = old_six_test
+
+
+class EnvironTests:
+
+    # grinning face emoji
+    UNICODE_EMOJI = six.u("\U0001F600")
+    BINARY_EMOJI = b"\xf0\x9f\x98\x80"
+
+    def test_set_environ(self, env_var):
+        six.environ['_six_test'] = UNICODE_EMOJI
+
+        assert six.environ['_six_test'] == UNICODE_EMOJI
+        assert six.environb['_six_test'] == BINARY_EMOJI
+        if PY3:
+            assert os.environ == UNICODE_EMOJI
+        else:
+            assert os.environ == BINARY_EMOJI
+
+    def test_set_environb(self, env_var):
+        six.environb['_six_test'] = BINARY_EMOJI
+
+        assert six.environ['_six_test'] == UNICODE_EMOJI
+        assert six.environb['_six_test'] == BINARY_EMOJI
+        if PY3:
+            assert os.environ == UNICODE_EMOJI
+        else:
+            assert os.environ == BINARY_EMOJI
+
+    def test_del_environ(self, env_var):
+        six.environ['_six_test'] = 'testing'
+        assert '_six_test' in os.environ
+        del six.environ['_six_test']
+        assert '_six_test' not in os.environ
+
+    def test_del_environb(self, env_var):
+        six.environb['_six_test'] = 'testing'
+        assert '_six_test' in os.environ
+        del six.environb['_six_test']
+        assert '_six_test' not in os.environ