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()