jupyter + pex entrypoint shim

jupyter_pex.py

from contextlib import contextmanager
import errno
import json
import os
import shutil
import sys
import tempfile

from notebook.notebookapp import main as notebook_main


def safe_mkdir(directory):
  try:
    os.makedirs(directory)
  except OSError as e:
    if e.errno != errno.EEXIST:
      raise


def safe_mkdir_for(path):
  safe_mkdir(os.path.dirname(path))


@contextmanager
def tmp_dir():
  tmp_path = tempfile.mkdtemp()
  try:
    yield tmp_path
  finally:
    shutil.rmtree(tmp_path)


class PexJupyterRunner(object):
  """A helper for running Jupyter notebooks from a pex.

  Currently, Jupyter notebook is made aware of various methods of executing new notebook
  "kernels" by way of on-disk configuration both in the Jupyter wheel as well as from local
  search paths on disk. The standard Python kernels do not work in the pex context because
  they assume e.g. a traditional python installation environment where the local site-packages
  is mutated with installs or where the surrounding environment is constructed via virtualenv
  tooling. Furthermore, the launch configuration is statically defined in JSON files on the
  filesystem and not in e.g. python code (jupyter-client.readthedocs.io/en/latest/kernels.html).

  An example of the default Python 2 kernel.json:

    {
     "display_name": "Python 2",
     "language": "python",
     "argv": ["python", "-m", "ipykernel_launcher", "-f", "{connection_file}"]
    }

  In order for this to work in a pex context, the launching needs to be self-referential with
  entrypoints controlled via environment variables. Because the configuration is in static JSON,
  we cannot reference things like sys.executable or sys.argv[0] in the static config in order
  to achieve self-reference.

  To hack around this for now, this shim is used as a surrogate for the notebook server entrypoint
  to perform the necessary configuration by writing static (but self-referential) JSON while in
  the pex execution context just prior to launching the notebook server.

  This hack works for the moment, but more formal follow-up is needed on the Github issue here:

      https://github.com/jupyter/notebook/issues/2636

  """

  @classmethod
  def get_config_path(cls, base_path):
    return os.path.join(base_path, 'kernels/pex/kernel.json')

  @staticmethod
  def get_pex_path():
    return os.path.abspath(sys.argv[0])

  @staticmethod
  def set_jupyter_path(path):
    os.environ['JUPYTER_PATH'] = str(path)

  @staticmethod
  def render_config(pex_path, display_name=None):
    display_name = display_name or 'PEX/{}'.format(os.path.basename(pex_path))
    return json.dumps(
      dict(
        display_name=display_name,
        env=dict(PEX_MODULE='ipykernel_launcher'),
        language='python',
        argv=[sys.executable, pex_path, '-f', '{connection_file}']
      )
    )

  @classmethod
  def write_config(cls, path):
    pex_path = cls.get_pex_path()
    config_path = cls.get_config_path(path)
    kernel_json = cls.render_config(pex_path)
    safe_mkdir_for(config_path)
    with open(config_path, 'wb') as f:
      f.write(kernel_json)

  @staticmethod
  def start_notebook_server():
    notebook_main()

  @classmethod
  def run(cls):
    with tmp_dir() as tmp:
      print('[shim] Writing pex kernel config')
      cls.write_config(tmp)

      print('[shim] Setting JUPYTER_PATH={}'.format(tmp))
      cls.set_jupyter_path(tmp)

      print('[shim] Launching notebook server')
      cls.start_notebook_server()


def launcher():
  PexJupyterRunner.run()

Great stuff, thanks for sharing! This helped me launch a Jupyter notebook for a Django project built with Pants.