Bug 1473498 - Fix Python 3 environment variables with subprocess r=gl…

…andium

On Windows in Python 2, the subprocess module requires the use of bytes with
the 'env' argument. For that reason, we would sometimes use byte strings with
'os.environ' like so:

    os.environ[b"FOO"] = b"bar"

However, this is a failure with Python 3 as 'os.environ' must only be used with
the text type. This patch creates a new 'setenv' helper that ensures we create
new environment with 'bytes' on Python 2, and 'text' on Python 3.

Differential Revision: https://phabricator.services.mozilla.com/D38237

build/mach_bootstrap.py

@@ -195,6 +195,7 @@ def bootstrap(topsrcdir, mozilla_dir=None):
                                              'build/virtualenv_packages.txt')]
     import mach.base
     import mach.main
+    from mach.util import setenv
     from mozboot.util import get_state_dir
 
     from mozbuild.util import patch_main
@@ -347,7 +348,7 @@ def populate_context(context, key=None):
     # Note which process is top-level so that recursive mach invocations can avoid writing
     # telemetry data.
     if 'MACH_MAIN_PID' not in os.environ:
-        os.environ[b'MACH_MAIN_PID'] = str(os.getpid()).encode('ascii')
+        setenv('MACH_MAIN_PID', str(os.getpid()))
 
     driver = mach.main.Mach(os.getcwd())
     driver.populate_context_handler = populate_context

python/mach/mach/main.py

@@ -36,6 +36,7 @@
 from .dispatcher import CommandAction
 from .logging import LoggingManager
 from .registrar import Registrar
+from .util import setenv
 
 SUGGEST_MACH_BUSTED = r'''
 You can invoke |./mach busted| to check if this issue is already on file. If it
@@ -354,7 +355,7 @@ def run(self, argv, stdin=None, stdout=None, stderr=None):
             # is a TTY. This provides a mechanism to allow said processes to
             # enable emitting code codes, for example.
             if os.isatty(orig_stdout.fileno()):
-                os.environ[b'MACH_STDOUT_ISATTY'] = b'1'
+                setenv('MACH_STDOUT_ISATTY', '1')
 
             return self._run(argv)
         except KeyboardInterrupt:

python/mach/mach/test/common.py

@@ -23,13 +23,14 @@
 class TestBase(unittest.TestCase):
     provider_dir = os.path.join(here, 'providers')
 
-    def get_mach(self, provider_file=None, entry_point=None, context_handler=None):
+    @classmethod
+    def get_mach(cls, provider_file=None, entry_point=None, context_handler=None):
         m = Mach(os.getcwd())
         m.define_category('testing', 'Mach unittest', 'Testing for mach core', 10)
         m.populate_context_handler = context_handler
 
         if provider_file:
-            m.load_commands_from_file(os.path.join(self.provider_dir, provider_file))
+            m.load_commands_from_file(os.path.join(cls.provider_dir, provider_file))
 
         if entry_point:
             m.load_commands_from_entry_point(entry_point)

python/mach/mach/test/python.ini

@@ -9,5 +9,6 @@ skip-if = python == 3
 [test_error_output.py]
 skip-if = python == 3
 [test_logger.py]
+[test_mach.py]
 [test_telemetry.py]
 skip-if = python == 3

python/mach/mach/test/test_mach.py

@@ -0,0 +1,34 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import absolute_import, unicode_literals
+
+import os
+
+from mach.test.common import TestBase
+from mozunit import main
+
+
+def test_set_isatty_environ(monkeypatch):
+    # Make sure the 'MACH_STDOUT_ISATTY' variable gets set.
+    monkeypatch.delenv('MACH_STDOUT_ISATTY', raising=False)
+    monkeypatch.setattr(os, 'isatty', lambda fd: True)
+
+    m = TestBase.get_mach()
+    orig_run = m._run
+    env_is_set = []
+
+    def wrap_run(*args, **kwargs):
+        env_is_set.append('MACH_STDOUT_ISATTY' in os.environ)
+        return orig_run(*args, **kwargs)
+
+    monkeypatch.setattr(m, '_run', wrap_run)
+
+    ret = m.run([])
+    assert ret == 0
+    assert env_is_set[0]
+
+
+if __name__ == '__main__':
+    main()

python/mach/mach/util.py

@@ -0,0 +1,28 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+from __future__ import absolute_import, unicode_literals
+
+import os
+import sys
+
+
+def setenv(key, value):
+    """Compatibility shim to ensure the proper string type is used with
+    os.environ for the version of Python being used.
+    """
+    encoding = "mbcs" if sys.platform == "win32" else "utf-8"
+
+    if sys.version_info[0] == 2:
+        if isinstance(key, unicode):
+            key = key.encode(encoding)
+        if isinstance(value, unicode):
+            value = value.encode(encoding)
+    else:
+        if isinstance(key, bytes):
+            key = key.decode(encoding)
+        if isinstance(value, bytes):
+            value = value.decode(encoding)
+
+    os.environ[key] = value