Merge pull request #472 from kevin-bates/templated-env

Add support for templated env entries

Author: blink1073

docs/kernels.rst

@@ -152,7 +152,10 @@ JSON serialised dictionary containing the following keys and values:
   the client will default to ``signal`` mode.
 - **env** (optional): A dictionary of environment variables to set for the kernel.
   These will be added to the current environment variables before the kernel is
-  started.
+  started.  Existing environment variables can be referenced using ``${<ENV_VAR>}`` and
+  will be substituted with the corresponding value.  Administrators should note that use
+  of ``${<ENV_VAR>}`` can expose sensitive variables and should use only in controlled
+  circumstances.
 - **metadata** (optional): A dictionary of additional attributes about this
   kernel; used by clients to aid in kernel selection. Metadata added
   here should be namespaced for the tool reading and writing that metadata.

jupyter_client/manager.py

@@ -241,19 +241,34 @@ def start_kernel(self, **kw):
         # If set, it can bork all the things.
         env.pop('PYTHONEXECUTABLE', None)
         if not self.kernel_cmd:
-            # If kernel_cmd has been set manually, don't refer to a kernel spec
-            # Environment variables from kernel spec are added to os.environ
-            env.update(self.kernel_spec.env or {})
+            # If kernel_cmd has been set manually, don't refer to a kernel spec.
+            # Environment variables from kernel spec are added to os.environ.
+            env.update(self._get_env_substitutions(self.kernel_spec.env, env))
         elif self.extra_env:
-            env.update(self.extra_env)
+            env.update(self._get_env_substitutions(self.extra_env, env))
 
         # launch the kernel subprocess
         self.log.debug("Starting kernel: %s", kernel_cmd)
-        self.kernel = self._launch_kernel(kernel_cmd, env=env,
-                                    **kw)
+        self.kernel = self._launch_kernel(kernel_cmd, env=env, **kw)
         self.start_restarter()
         self._connect_control_socket()
 
+    def _get_env_substitutions(self, templated_env, substitution_values):
+        """ Walks env entries in templated_env and applies possible substitutions from current env
+            (represented by substitution_values).
+            Returns the substituted list of env entries.
+        """
+        substituted_env = {}
+        if templated_env:
+            from string import Template
+
+            # For each templated env entry, fill any templated references
+            # matching names of env variables with those values and build
+            # new dict with substitutions.
+            for k, v in templated_env.items():
+                substituted_env.update({k: Template(v).safe_substitute(substitution_values)})
+        return substituted_env
+
     def request_shutdown(self, restart=False):
         """Send a shutdown request via control channel
         """

jupyter_client/tests/signalkernel.py

@@ -4,6 +4,7 @@
 # Distributed under the terms of the Modified BSD License.
 
 from __future__ import print_function
+import os
 
 from subprocess import Popen, PIPE
 import sys
@@ -38,6 +39,8 @@ def do_execute(self, code, silent, store_history=True, user_expressions=None,
             reply['user_expressions']['pid'] = self.children[-1].pid
         elif code == 'check':
             reply['user_expressions']['poll'] = [ child.poll() for child in self.children ]
+        elif code == 'env':
+            reply['user_expressions']['env'] = os.getenv("TEST_VARS", "")
         elif code == 'sleep':
             try:
                 time.sleep(10)

jupyter_client/tests/test_kernelmanager.py

@@ -41,6 +41,7 @@ def _install_test_kernel(self):
                          '-m', 'jupyter_client.tests.signalkernel',
                          '-f', '{connection_file}'],
                 'display_name': "Signal Test Kernel",
+                'env': {'TEST_VARS': '${TEST_VARS}:test_var_2'},
             }))
 
     def _get_tcp_km(self):
@@ -131,6 +132,63 @@ def test_start_new_kernel(self):
         self.assertTrue(km.is_alive())
         self.assertTrue(kc.is_alive())
 
+    def _env_test_body(self, kc):
+
+        def execute(cmd):
+            kc.execute(cmd)
+            reply = kc.get_shell_msg(TIMEOUT)
+            content = reply['content']
+            self.assertEqual(content['status'], 'ok')
+            return content
+
+        reply = execute('env')
+        self.assertIsNotNone(reply)
+        self.assertEquals(reply['user_expressions']['env'], 'test_var_1:test_var_2')
+
+    def test_templated_kspec_env(self):
+        self._install_test_kernel()
+        km, kc = start_new_kernel(kernel_name='signaltest')
+        self.addCleanup(kc.stop_channels)
+        self.addCleanup(km.shutdown_kernel)
+
+        self.assertTrue(km.is_alive())
+        self.assertTrue(kc.is_alive())
+
+        self._env_test_body(kc)
+
+    def _start_kernel_with_cmd(self, kernel_cmd, extra_env, **kwargs):
+        """Start a new kernel, and return its Manager and Client"""
+        km = KernelManager(kernel_name='signaltest')
+        km.kernel_cmd = kernel_cmd
+        km.extra_env = extra_env
+        km.start_kernel(**kwargs)
+        kc = km.client()
+        kc.start_channels()
+        try:
+            kc.wait_for_ready(timeout=60)
+        except RuntimeError:
+            kc.stop_channels()
+            km.shutdown_kernel()
+            raise
+
+        return km, kc
+
+    def test_templated_extra_env(self):
+        self._install_test_kernel()
+        kernel_cmd = [sys.executable,
+                         '-m', 'jupyter_client.tests.signalkernel',
+                         '-f', '{connection_file}']
+        extra_env = {'TEST_VARS': '${TEST_VARS}:test_var_2'}
+
+        km, kc = self._start_kernel_with_cmd(kernel_cmd, extra_env)
+        self.addCleanup(kc.stop_channels)
+        self.addCleanup(km.shutdown_kernel)
+
+        self.assertTrue(km.is_alive())
+        self.assertTrue(kc.is_alive())
+
+        self._env_test_body(kc)
+
 
 class TestParallel:
 

jupyter_client/tests/utils.py

@@ -29,6 +29,7 @@ def start(self):
             'JUPYTER_DATA_DIR': pjoin(td.name, 'jupyter_data'),
             'JUPYTER_RUNTIME_DIR': pjoin(td.name, 'jupyter_runtime'),
             'IPYTHONDIR': pjoin(td.name, 'ipython'),
+            'TEST_VARS': 'test_var_1',
         })
         self.env_patch.start()