Merge pull request #541 from fesaille/cli

Add env parameter to kernel installation

ipykernel/kernelspec.py

@@ -82,7 +82,7 @@ def write_kernel_spec(path=None, overrides=None, extra_arguments=None):
 
 
 def install(kernel_spec_manager=None, user=False, kernel_name=KERNEL_NAME, display_name=None,
-            prefix=None, profile=None):
+            prefix=None, profile=None, env=None):
     """Install the IPython kernelspec for Jupyter
     
     Parameters
@@ -103,6 +103,10 @@ def install(kernel_spec_manager=None, user=False, kernel_name=KERNEL_NAME, displ
     prefix: str, optional
         Specify an install prefix for the kernelspec.
         This is needed to install into a non-default location, such as a conda/virtual-env.
+    env: dict, optional
+        A dictionary of extra environment variables for the kernel.
+        These will be added to the current environment variables before the
+        kernel is started
 
     Returns
     -------
@@ -126,6 +130,8 @@ def install(kernel_spec_manager=None, user=False, kernel_name=KERNEL_NAME, displ
             overrides["display_name"] = 'Python %i [profile=%s]' % (sys.version_info[0], profile)
     else:
         extra_arguments = None
+    if env:
+        overrides['env'] = env
     path = write_kernel_spec(overrides=overrides, extra_arguments=extra_arguments)
     dest = kernel_spec_manager.install_kernel_spec(
         path, kernel_name=kernel_name, user=user, prefix=prefix)
@@ -168,10 +174,14 @@ def start(self):
         parser.add_argument('--sys-prefix', action='store_const', const=sys.prefix, dest='prefix',
             help="Install to Python's sys.prefix."
             " Shorthand for --prefix='%s'. For use in conda/virtual-envs." % sys.prefix)
+        parser.add_argument('--env', action='append', nargs=2, metavar=('ENV', 'VALUE'),
+            help="Set environment variables for the kernel.")
         opts = parser.parse_args(self.argv)
+        if opts.env:
+            opts.env = {k:v for (k, v) in opts.env}
         try:
             dest = install(user=opts.user, kernel_name=opts.name, profile=opts.profile,
-                           prefix=opts.prefix, display_name=opts.display_name)
+                           prefix=opts.prefix, display_name=opts.display_name, env=opts.env)
         except OSError as e:
             if e.errno == errno.EACCES:
                 print(e, file=sys.stderr)

ipykernel/tests/test_kernelspec.py

@@ -21,6 +21,8 @@
     RESOURCES,
 )
 
+import pytest
+
 import nose.tools as nt
 
 pjoin = os.path.join
@@ -140,3 +142,27 @@ def test_install_display_name_overrides_profile():
     with open(spec) as f:
         spec = json.load(f)
     assert spec["display_name"] == "Display"
+
+
+@pytest.mark.parametrize("env", [
+    None,
+    dict(spam="spam"),
+    dict(spam="spam", foo='bar')
+])
+def test_install_env(tmp_path, env):
+    # python 3.5 // tmp_path must be converted to str
+    with mock.patch('jupyter_client.kernelspec.SYSTEM_JUPYTER_PATH',
+            [str(tmp_path)]):
+        install(env=env)
+
+    spec = tmp_path / 'kernels' / KERNEL_NAME / "kernel.json"
+    with spec.open() as f:
+        spec = json.load(f)
+
+    if env:
+        assert len(env) == len(spec['env'])
+        for k, v in env.items():
+            assert spec['env'][k] == v
+    else:
+        assert 'env' not in spec
+