adds --conda-activate option when installing kernel

Author: edublancas

bash_kernel/assets/bashrc-conda

@@ -0,0 +1,21 @@
+# Different platforms have different names for the systemwide bashrc
+if [[ -f /etc/bashrc ]]; then
+  source /etc/bashrc
+fi
+if [[ -f /etc/bash.bashrc ]]; then
+  source /etc/bash.bashrc
+fi
+if [[ -f ~/.bashrc ]]; then
+  source ~/.bashrc
+fi
+
+# Reset PS1 so pexpect can find it
+PS1="$"
+
+# Unset PROMPT_COMMAND, so that it can't change PS1 to something unexpected.
+unset PROMPT_COMMAND
+
+
+eval "$(conda shell.bash hook)"
+
+conda activate {CONDA_DEFAULT_ENV}

bash_kernel/install.py

@@ -13,11 +13,28 @@
  "env":{"PS1": "$"}
 }
 
-def install_my_kernel_spec(user=True, prefix=None):
+def install_my_kernel_spec(user=True, prefix=None, conda_activate=False):
+    if conda_activate:
+        CONDA_DEFAULT_ENV = os.environ.get("CONDA_DEFAULT_ENV")
+
+        if not CONDA_DEFAULT_ENV:
+            raise RuntimeError(
+                "This script must be run from a conda environment "
+                "when using the --conda-activate flag"
+            )
+
+        settings = {"conda_default_env": CONDA_DEFAULT_ENV}
+    else:
+        settings = {"conda_default_env": None}
+
     with TemporaryDirectory() as td:
         os.chmod(td, 0o755) # Starts off as 700, not user readable
         with open(os.path.join(td, 'kernel.json'), 'w') as f:
             json.dump(kernel_json, f, sort_keys=True)
+
+        with open(os.path.join(td, 'settings.json'), 'w') as f:
+            json.dump(settings, f, sort_keys=True)
+
         # TODO: Copy resources once they're specified
 
         print('Installing IPython kernel spec')
@@ -51,6 +68,12 @@ def main(argv=None):
         help='Install KernelSpec in this prefix',
         default=None
     )
+    parser.add_argument(
+        '--conda-activate',
+        help='Initialize the bash kernel with the current conda environment',
+        action='store_true',
+        dest='conda_activate'
+    )
 
     args = parser.parse_args(argv)
 
@@ -63,7 +86,7 @@ def main(argv=None):
     elif args.user or not _is_root():
         user = True
 
-    install_my_kernel_spec(user=user, prefix=prefix)
+    install_my_kernel_spec(user=user, prefix=prefix, conda_activate=args.conda_activate)
 
 if __name__ == '__main__':
     main()

bash_kernel/kernel.py

@@ -1,13 +1,25 @@
+import json
+import tempfile
 from ipykernel.kernelbase import Kernel
+from jupyter_client.kernelspec import find_kernel_specs
 from pexpect import replwrap, EOF
 import pexpect
+from pathlib import Path
 
 from subprocess import check_output
 import os.path
 import uuid
 import random
 import string
 
+try:
+    import bash_kernel
+# installation reads __version__. since the package isn't ready yet
+# this will fail, so we add the except clause
+except ModuleNotFoundError:
+    bash_kernel = None
+
+
 import re
 import signal
 
@@ -101,9 +113,42 @@ def __init__(self, **kwargs):
         rand = ''.join(random.choices(string.ascii_uppercase, k=12))
         self.unique_prompt = "PROMPT_" + rand
         Kernel.__init__(self, **kwargs)
-        self._start_bash()
+
+        # determine which bashrc to use
+        kernelspecs = find_kernel_specs()
+        kernelspec_bash = kernelspecs['bash']
+
+        if not kernelspec_bash:
+            raise RuntimeError("Can't find bash kernelspec")
+
+        settings = json.loads(Path(kernelspec_bash, 'settings.json').read_text())
+        CONDA_DEFAULT_ENV = settings.get("conda_default_env")
+
+        if CONDA_DEFAULT_ENV:
+            print("Generating temporary bashrc...")
+            template = os.path.join(os.path.dirname(bash_kernel.__file__), 'assets',
+                                    'bashrc-conda')
+            template_content = (Path(template).read_text()
+                                .format(CONDA_DEFAULT_ENV=CONDA_DEFAULT_ENV))
+
+            self._bashrc_tmp = tempfile.NamedTemporaryFile(mode='w')
+            self._bashrc_tmp.write(template_content)
+            self._bashrc_tmp.flush()
+            self._path_to_bashrc = self._bashrc_tmp.name
+            self._start_bash()
+        else:
+            bashrc = os.path.join(os.path.dirname(pexpect.__file__), 'bashrc.sh')
+            self._bashrc_tmp = None
+            self._path_to_bashrc = bashrc
+            self._start_bash()
+
+
         self._known_display_ids = set()
 
+    def __del__(self):
+        if self._bashrc_tmp:
+            self._bashrc_tmp.close()
+
     def _start_bash(self):
         # Signal handlers are inherited by forked processes, and we can't easily
         # reset it from the subprocess. Since kernelapp ignores SIGINT except in
@@ -115,8 +160,7 @@ def _start_bash(self):
             # bash() function of pexpect/replwrap.py.  Look at the
             # source code there for comments and context for
             # understanding the code here.
-            bashrc = os.path.join(os.path.dirname(pexpect.__file__), 'bashrc.sh')
-            child = pexpect.spawn("bash", ['--rcfile', bashrc], echo=False,
+            child = pexpect.spawn("bash", ['--rcfile', self._path_to_bashrc], echo=False,
                                   encoding='utf-8', codec_errors='replace')
             # Following comment stolen from upstream's REPLWrap:
             # If the user runs 'env', the value of PS1 will be in the output. To avoid