Force environment variables to str with Python2 on Windows (theskumar#101)

greyli · theskumar · commit 43af2c57a362 · 2018-10-31T16:55:54.000+05:30
* Force environment variables to str

With Python2 on Windows, force environment variables to str to avoid "TypeError: environment can only contain strings " in Python's subprocess.py.

* PEP8 fix

* Add Python3 compatible

* Fix flake8 fails with F821

* Fix flake8 fails with F821
diff --git a/dotenv/compat.py b/dotenv/compat.py
@@ -1,4 +1,9 @@
+import sys
 try:
     from StringIO import StringIO  # noqa
 except ImportError:
     from io import StringIO  # noqa
+
+PY2 = sys.version_info[0] == 2
+WIN = sys.platform.startswith('win')
+text_type = unicode if PY2 else str  # noqa
diff --git a/dotenv/main.py b/dotenv/main.py
@@ -11,7 +11,7 @@
 import warnings
 from collections import OrderedDict
 
-from .compat import StringIO
+from .compat import StringIO, PY2, WIN, text_type
 
 __escape_decoder = codecs.getdecoder('unicode_escape')
 __posix_variable = re.compile('\$\{[^\}]*\}')  # noqa
@@ -95,6 +95,12 @@ def set_as_environment_variables(self, override=False):
         for k, v in self.dict().items():
             if k in os.environ and not override:
                 continue
+            # With Python2 on Windows, force environment variables to str to avoid
+            # "TypeError: environment can only contain strings" in Python's subprocess.py.
+            if PY2 and WIN:
+                if isinstance(k, text_type) or isinstance(v, text_type):
+                    k = k.encode('ascii')
+                    v = v.encode('ascii')
             os.environ[k] = v
 
         return True
