From f977be2868b14a4028149a78603c208377689504 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Tue, 17 Jul 2018 23:16:54 +0100
Subject: [PATCH] issue #291: permit supplying a full Python argv.

---
 docs/api.rst        | 10 +++++++---
 mitogen/parent.py   | 17 +++++++++++++++--
 tests/local_test.py | 39 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/docs/api.rst b/docs/api.rst
index bc88a622..6efca6dd 100644
--- a/docs/api.rst
+++ b/docs/api.rst
@@ -483,9 +483,13 @@ Router Class
             determine its installation prefix. This is required to support
             virtualenv.
 
-        :param str python_path:
-            Path to the Python interpreter to use for bootstrap. Defaults to
-            :data:`sys.executable`. For SSH, defaults to ``python``.
+        :param str|list python_path:
+            String or list path to the Python interpreter to use for bootstrap.
+            Defaults to :data:`sys.executable` for local connections, and
+            ``python`` for remote connections.
+
+            It is possible to pass a list to invoke Python wrapped using
+            another tool, such as ``["/usr/bin/env", "python"]``.
 
         :param bool debug:
             If :data:`True`, arrange for debug logging (:py:meth:`enable_debug`) to
diff --git a/mitogen/parent.py b/mitogen/parent.py
index 38484874..21332454 100644
--- a/mitogen/parent.py
+++ b/mitogen/parent.py
@@ -871,6 +871,19 @@ class Stream(mitogen.core.Stream):
         fp.close()
         os.write(1,'MITO001\n'.encode())
 
+    def get_python_argv(self):
+        """
+        Return the initial argument vector elements necessary to invoke Python,
+        by returning a 1-element list containing :attr:`python_path` if it is a
+        string, or simply returning it if it is already a list.
+
+        This allows emulation of existing tools where the Python invocation may
+        be set to e.g. `['/usr/bin/env', 'python']`.
+        """
+        if isinstance(self.python_path, list):
+            return self.python_path
+        return [self.python_path]
+
     def get_boot_command(self):
         source = inspect.getsource(self._first_stage)
         source = textwrap.dedent('\n'.join(source.strip().split('\n')[2:]))
@@ -886,8 +899,8 @@ class Stream(mitogen.core.Stream):
         # codecs.decode() requires a bytes object. Since we must be compatible
         # with 2.4 (no bytes literal), an extra .encode() either returns the
         # same str (2.x) or an equivalent bytes (3.x).
-        return [
-            self.python_path, '-c',
+        return self.get_python_argv() + [
+            '-c',
             'import codecs,os,sys;_=codecs.decode;'
             'exec(_(_("%s".encode(),"base64"),"zip"))' % (encoded.decode(),)
         ]
diff --git a/tests/local_test.py b/tests/local_test.py
index 8ba248da..fbf5c1c8 100644
--- a/tests/local_test.py
+++ b/tests/local_test.py
@@ -1,5 +1,6 @@
 
 import os
+import sys
 
 import unittest2
 
@@ -11,6 +12,14 @@ import testlib
 import plain_old_module
 
 
+def get_sys_executable():
+    return sys.executable
+
+
+def get_os_environ():
+    return dict(os.environ)
+
+
 class LocalTest(testlib.RouterMixin, unittest2.TestCase):
     stream_class = mitogen.ssh.Stream
 
@@ -20,5 +29,35 @@ class LocalTest(testlib.RouterMixin, unittest2.TestCase):
         self.assertEquals('local.%d' % (pid,), context.name)
 
 
+class PythonPathTest(testlib.RouterMixin, unittest2.TestCase):
+    stream_class = mitogen.ssh.Stream
+
+    def test_inherited(self):
+        context = self.router.local()
+        self.assertEquals(sys.executable, context.call(get_sys_executable))
+
+    def test_string(self):
+        os.environ['PYTHON'] = sys.executable
+        context = self.router.local(
+            python_path=testlib.data_path('env_wrapper.sh'),
+        )
+        self.assertEquals(sys.executable, context.call(get_sys_executable))
+        env = context.call(get_os_environ)
+        self.assertEquals('1', env['EXECUTED_VIA_ENV_WRAPPER'])
+
+    def test_list(self):
+        context = self.router.local(
+            python_path=[
+                testlib.data_path('env_wrapper.sh'),
+                "magic_first_arg",
+                sys.executable
+            ]
+        )
+        self.assertEquals(sys.executable, context.call(get_sys_executable))
+        env = context.call(get_os_environ)
+        self.assertEquals('magic_first_arg', env['ENV_WRAPPER_FIRST_ARG'])
+        self.assertEquals('1', env['EXECUTED_VIA_ENV_WRAPPER'])
+
+
 if __name__ == '__main__':
     unittest2.main()