Adding proper activation for conda envs to ipykernel creation

ipykernel/kernelspec.py

@@ -59,37 +59,47 @@ def get_kernel_dict(extra_arguments=None):
     }
 
 
-def write_kernel_spec(path=None, overrides=None, extra_arguments=None):
+def write_kernel_spec(path=None, overrides=None, extra_arguments=None, conda_env=False):
     """Write a kernel spec directory to `path`
-    
+
     If `path` is not specified, a temporary directory is created.
     If `overrides` is given, the kernelspec JSON is updated before writing.
-    
+
     The path to the kernelspec is always returned.
     """
     if path is None:
         path = os.path.join(tempfile.mkdtemp(suffix='_kernels'), KERNEL_NAME)
-    
+
     # stage resources
     shutil.copytree(RESOURCES, path)
     # write kernel.json
     kernel_dict = get_kernel_dict(extra_arguments)
+    if conda_env:
+        # argv is the list that jupyter invokes via a subprocess call
+        argv = kernel_dict.get('argv', [])
+        # The first element of argv is something like /home/user/miniconda/envs/<conda env name>/bin/python
+        # We want to pull out the <conda env name> so we can pass it to conda run. It's the third-from-last
+        # element in the path, so we'll get that with the os-specific path split
+        conda_env_name = argv[0].split(os.sep)[-3]
+        # Format the new argv and put it back in the original dict as the argv key
+        new_argv = ['conda', 'run', '-n', conda_env_name, 'python'] + argv[1:]
+        kernel_dict['argv'] = new_argv
 
     if overrides:
         kernel_dict.update(overrides)
     with open(pjoin(path, 'kernel.json'), 'w') as f:
         json.dump(kernel_dict, f, indent=1)
-    
+
     return path
 
 
 def install(kernel_spec_manager=None, user=False, kernel_name=KERNEL_NAME, display_name=None,
-            prefix=None, profile=None):
+            prefix=None, profile=None, conda_env=False):
     """Install the IPython kernelspec for Jupyter
-    
+
     Parameters
     ----------
-    
+
     kernel_spec_manager: KernelSpecManager [optional]
         A KernelSpecManager to use for installation.
         If none provided, a default instance will be created.
@@ -105,10 +115,12 @@ 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.
+    conda_env: bool, optional
+        Format the kernelspec so that it properly activates the conda environment, if your kernel
+        is running the python executable out of a conda environment.
 
     Returns
     -------
-
     The path where the kernelspec was installed.
     """
     if kernel_spec_manager is None:
@@ -128,7 +140,7 @@ 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
-    path = write_kernel_spec(overrides=overrides, extra_arguments=extra_arguments)
+    path = write_kernel_spec(overrides=overrides, extra_arguments=extra_arguments, conda_env=conda_env)
     dest = kernel_spec_manager.install_kernel_spec(
         path, kernel_name=kernel_name, user=user, prefix=prefix)
     # cleanup afterward
@@ -143,12 +155,12 @@ def install(kernel_spec_manager=None, user=False, kernel_name=KERNEL_NAME, displ
 class InstallIPythonKernelSpecApp(Application):
     """Dummy app wrapping argparse"""
     name = 'ipython-kernel-install'
-    
+
     def initialize(self, argv=None):
         if argv is None:
             argv = sys.argv[1:]
         self.argv = argv
-    
+
     def start(self):
         import argparse
         parser = argparse.ArgumentParser(prog=self.name,
@@ -170,10 +182,23 @@ 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('--conda', action='store_true', dest='conda_env',
+            default=False,
+            help="Ensure that the conda environment is properly activated for kernel launch")
+        # Temporary addition, will be removed after dev is finished
+        parser.add_argument('--pdb', action="store_true", help="Enable PDB debugging on exception",
+                            default=False)
         opts = parser.parse_args(self.argv)
+        if opts.pdb:
+            import pdb
+            # set the pdb_hook as the except hook for all exceptions
+            def pdb_hook(exctype, value, traceback):
+                pdb.post_mortem(traceback)
+            sys.excepthook = pdb_hook
+
         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, conda_env=opts.conda_env)
         except OSError as e:
             if e.errno == errno.EACCES:
                 print(e, file=sys.stderr)