Repository: jupyter/jupyter_client
Pull request #472 (2019-08-31 22:10:47)

Add support for templated env entries
I've encountered a number of people (including myself) that would like templated environment variable entries in kernelspecs, so I thought I'd implement it add see if this is something useful to the community.

With this change, kernelspec env entries (or even those clients that set `extra_env` on the `KernelManager` instance (along with `kernel_cmd`)) can reference existing env variables using `${ENV_VAR}` within various `env` entries and expect those substitutions to be performed prior to the kernel's launch.

Here's an example kernelspec env entry:

```json
"env": { 
    "PYTHONPATH": "my_special_python_path:${PYTHONPATH}", 
    "FOO": "BAR", 
    "MY_APP_PATH": "my_bin:${PATH}"
}
```
Following the kernel's invocation, it's environment will include the following variable/values:
```
PYTHONPATH=my_special_python_path:<notebook server's PYTHONPATH>
FOO=BAR
MY_APP_PATH=my_bin:<notebook server's PATH>
```
If the templated env variable is not in the launching environment, the templated value will remain in place - just as it shows in the kernelspec's `kernel.json` file.

Kernelspec administrators should be cognizant of potential to share environments.  However, given that's case currently, this shouldn't introduce any new issues.


Patch:
diff --git a/docs/kernels.rst b/docs/kernels.rst
index 41e2eb8..f3c06d1 100644
--- a/docs/kernels.rst
+++ b/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.
diff --git a/jupyter_client/manager.py b/jupyter_client/manager.py
index 4e13ffd..1d47871 100644
--- a/jupyter_client/manager.py
+++ b/jupyter_client/manager.py
@@ -241,19 +241,34 @@ class KernelManager(ConnectionFileMixin):
         # 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
         """


Test patch:
diff --git a/jupyter_client/tests/signalkernel.py b/jupyter_client/tests/signalkernel.py
index e56ea44..44a53ee 100644
--- a/jupyter_client/tests/signalkernel.py
+++ b/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 @@ class SignalTestKernel(Kernel):
             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)
diff --git a/jupyter_client/tests/test_kernelmanager.py b/jupyter_client/tests/test_kernelmanager.py
index aee63a9..58010bc 100644
--- a/jupyter_client/tests/test_kernelmanager.py
+++ b/jupyter_client/tests/test_kernelmanager.py
@@ -41,6 +41,7 @@ class TestKernelManager(TestCase):
                          '-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 @@ class TestKernelManager(TestCase):
         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:
 
diff --git a/jupyter_client/tests/utils.py b/jupyter_client/tests/utils.py
index 505084a..64ce45d 100644
--- a/jupyter_client/tests/utils.py
+++ b/jupyter_client/tests/utils.py
@@ -29,6 +29,7 @@ class test_env(object):
             '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()