CLI Tool to Manage Conda Envs & Jupyter Notebooks

May 29, 2021

CondaShell-Scripting

I’ve been using anaconda environments and jupyter notebook a lot recently. One thing I love about conda envs are that once created it can be used anywhere. Python’s venv requires you to set up the virtual env in a directory. My anaconda setup get too messy at times with too many envs. Sometimes I forget which env I used to do the work and ends up creating a new one. Also the conda commands seems pretty hard to remember (I’m getting used to it now). So the best possible solution that I could come up with was to write a script.

The script let’s me start the jupyter notebook from any directory on an existing environment or a new one. It lists all the existing environments on the screen and I can choose from it. This script has been a sigh of relief that I could create conda envs and start notebooks easily.

Working on the script

I started with a python script rather than a shell script so that I could use python libraries to make the work easier. But I ran into many issues along the way, so I had to use both python and shell scripts to make it work.

Dependencies

The only important dependency here is a python library called PyInquirer. It lets us build an interactive CLI with different types of prompts. You should also install notebook as a global package, if you don’t want to do that, installing it in the current env would also work.

  1. Create a python virtual env in a directory

    python3 -m venv env
source env/bin/activate

  2. Install dependencies

    pip install PyInquirer
pip install notebook

How it works?

It first asks us if we want to create a new env. If yes, then we should enter the env name and dependencies. After that we can choose an env from conda env list. It then starts the selected environement. Also, it asks us if we want to open the notebook.

Coding it

  1. Create a python file cli.py and import dependencies

    # cli.py

from PyInquirer import prompt
import subprocess
import re, os, sys

# current working directory
CURR_WD = os.getcwd()

  2. Add prompts for creating a new env

    # cli.py

# prompts user whether to create new env
# returns a dictionary of answers
answers = prompt([
  {
      'type': 'confirm',
      'name': 'new',
      'message': 'Create a new environment?',
      'default': False
  }
])

new_env = answers['new']

    It prompts a yes or no question and stores the our reponse in a new variable.

  3. Create virtual env if new is True:

    # cli.py

if new_env is True:
   name = prompt([
       {
           'type': 'input',
           'name': 'name',
           'message': 'Name of the environment:'
       }
   ])['name']

   packages = prompt([
       {
           'type': 'input',
           'name': 'packages',
           'message': 'Enter the packages to be installed:'
       }
   ])['packages']

   add_notebook = prompt([
       {
           'type': 'confirm',
           'name': 'add_notebook',
           'message': 'Want to install jupyter notebook?',
           'default': False
       }
   ])['add_notebook']

   add_notebook = 1 if add_notebook else 0 
   
   # runs `create.sh` file with args
   subprocess.run('{}/create.sh {} {} {}'.format(CURR_WD, name, add_notebook, packages), shell=True)

    It asks users for info like name of env, packages to be installed and whether jupyter notebook to be installed. The subprocess.run function is used to run a shell command from python. Here it runs another shell script which is created in same directory as cli.py. The create.sh script is used to create a virtual env. The name, add_notebook and packages variables become the arguments to the script. It uses these arguments to create a conda env.

    Note: You might think why not just run the conda create command using subprocess, but I felt it would be easier to run the bash commands on a shell script.

  4. The create.sh file

    # create.sh

#!/bin/bash

. ~/anaconda3/etc/profile.d/conda.sh

# array of arguments
args=($@)

# dependencies
dps=''

# loop and create a string of dependencies
for ((i=2; i<$#; i++))
do
   dps+="${args[$i]} "
done

# create a conda env
conda create -y -n $1 $dps;

# activate the env
conda activate $1

# install notebook if `add_notebook` is 1
if [ $2 -eq "1" ]
then 
   pip install notebook
fi

# deactivate the env
conda deactivate

    It creates a string of the dependencies and run the conda create command. It then install jupyter notebook if add_notebook argument val was 1.

  5. Prompt for selecting conda env

    # cli.py

# runs `conda info --envs` and pipe the response to `screen_output`
screen_output = subprocess.run('conda info --envs', stdout=subprocess.PIPE, text=True, shell=True).stdout

# uses regex to find the env names from `screen_output`
envs_dirty = re.findall(r'\n(\w*)', screen_output)

# removes empty values
envs_clean = [env for env in envs_dirty if env]

# Prompts the user to select an env
answers = prompt([
   {
       'type': 'list',
       'name': 'env',
       'message': 'Choose an environment:',
       'choices': envs_clean
   }
])

# env to be activated
env = answers['env']

  6. Prompt the user to choose whether to open a notebook in the current env

    # cli.py

answers = prompt([
   {
       'type': 'confirm',
       'name': 'open_notebook',
       'message': 'Wanna open jupyter notebook?',
       'default': False
   }
])

open_notebook = answers['open_notebook']
open_notebook = 1 if open_notebook else 0

  7. Creating a list of all active notebooks

    # cli.py

screen_output = subprocess.run('{}/venv/bin/jupyter notebook list'.format(CURR_WD),  stdout=subprocess.PIPE, text=True, shell=True).stdout

# creates a list of port number of active notebooks using regex
prev_hosts = re.findall(r'localhost:(\d*)', screen_output)

  8. Starting the environment and opening notebook

    # cli.py

try:
   subprocess.run('sh {}/start.sh {} {}'.format(CURR_WD, env, open_notebook), shell=True)
except KeyboardInterrupt: 
   screen_output = subprocess.run('{}/venv/bin/jupyter notebook list'.format(CURR_WD),  stdout=subprocess.PIPE, text=True, shell=True).stdout
   
   # create a list of all the running notebooks
   curr_hosts = re.findall(r'localhost:(\d*)', screen_output)
   
   # difference gives the list of notebooks started by start.sh
   hosts = list(set(curr_hosts) - set(prev_hosts))
   
   # stop the notebooks
   for host in hosts:
       res = subprocess.run('{}/venv/bin/jupyter notebook stop {}'.format(CURR_WD, host),  stdout=subprocess.PIPE, text=True, shell=True)
   
   # exit the program
   sys.exit()

    It calls the start.sh script which creates a virtual env and starts the notebook. The best way to close the notebook is to hit ctrl+c. Here the catch block captures the KeyboardInterrupt error and closes the notebook.

  9. The start.sh file

    # start.sh

#!/bin/sh

# run conda.sh
# . is an alternative to `source`
. ~/anaconda3/etc/profile.d/conda.sh

# activate conda env
conda activate $1

# start notebook
if [ $2 -eq "1" ]
then 
   jupyter notebook
fi

Last steps

Our script is now complete. The final steps are to make the make the script runnable from any directory. To do that we’ll give executable permissions to all the scripts and add it to .bashrc

chmod +x ./cli.py
chmod +x ./create.sh
chmod +x ./start.sh

Finally we can give an alias to the cli.py. To do that, open .bashrc vim ~/.bashrc and add an alias to the file

# .bashrc

alias open-notebook='<path_to_the_file>/cli.py'

Fire up a new terminal and we can run the script from anywhere using open-notebook command.

Final ‘cli.py’ file

from PyInquirer import prompt
import subprocess
import re, os, sys

# current working directory
CURR_WD = os.getcwd()

# prompts user whether to create new env
# returns a dictionary of answers
answers = prompt([
   {
       'type': 'confirm',
       'name': 'new',
       'message': 'Create a new environment?',
       'default': False
   }
])

new_env = answers['new']

if new_env is True:
    name = prompt([
        {
            'type': 'input',
            'name': 'name',
            'message': 'Name of the environment:'
        }
    ])['name']

    packages = prompt([
        {
            'type': 'input',
            'name': 'packages',
            'message': 'Enter the packages to be installed:'
        }
    ])['packages']

    add_notebook = prompt([
        {
            'type': 'confirm',
            'name': 'add_notebook',
            'message': 'Want to install jupyter notebook?',
            'default': False
        }
    ])['add_notebook']

    add_notebook = 1 if add_notebook else 0 
    
    # runs `create.sh` file with args
    subprocess.run('{}/create.sh {} {} {}'.format(CURR_WD, name, add_notebook, packages), shell=True)

# runs `conda info --envs` and pipe the response to `screen_output`
screen_output = subprocess.run('conda info --envs', stdout=subprocess.PIPE, text=True, shell=True).stdout

# uses regex to find the env names from `screen_output`
envs_dirty = re.findall(r'\n(\w*)', screen_output)

# removes empty values
envs_clean = [env for env in envs_dirty if env]

# Prompts the user to select an env
answers = prompt([
    {
        'type': 'list',
        'name': 'env',
        'message': 'Choose an environment:',
        'choices': envs_clean
    }
])

# env to be activated
env = answers['env']

answers = prompt([
    {
        'type': 'confirm',
        'name': 'open_notebook',
        'message': 'Wanna open jupyter notebook?',
        'default': False
    }
])

open_notebook = answers['open_notebook']
open_notebook = 1 if open_notebook else 0

screen_output = subprocess.run('{}/venv/bin/jupyter notebook list'.format(CURR_WD),  stdout=subprocess.PIPE, text=True, shell=True).stdout

# creates a list of port number of active notebooks using regex
prev_hosts = re.findall(r'localhost:(\d*)', screen_output)

try:
    subprocess.run('sh {}/start.sh {} {}'.format(CURR_WD, env, open_notebook), shell=True)
except KeyboardInterrupt: 
    screen_output = subprocess.run('{}/venv/bin/jupyter notebook list'.format(CURR_WD),  stdout=subprocess.PIPE, text=True, shell=True).stdout
    
    # create a list of all the running notebooks
    curr_hosts = re.findall(r'localhost:(\d*)', screen_output)
    
    # difference gives the list of notebooks started by start.sh
    hosts = list(set(curr_hosts) - set(prev_hosts))
    
    # stop the notebooks
    for host in hosts:
        res = subprocess.run('{}/venv/bin/jupyter notebook stop {}'.format(CURR_WD, host),  stdout=subprocess.PIPE, text=True, shell=True)
    
    # exit the program
    sys.exit()
Profile picture

Written by Roshan R who lives in India, building useful things. You should follow them on Twitter