Skip to content

Instantly share code, notes, and snippets.

@rcuza
Last active October 13, 2020 01:57
Show Gist options
  • Save rcuza/d9ec895ef0abe21036ffd58aeec96c40 to your computer and use it in GitHub Desktop.
Save rcuza/d9ec895ef0abe21036ffd58aeec96c40 to your computer and use it in GitHub Desktop.
Pants Plugin for Jupyter Repl for v1 Engine. Works with v1.26.

This code is shared under the Apache License, copyright Chartbeat.

Jupyter Pants Plugin

The organization of the plugin directory is as follows:

jupyter
├── BUILD
├── __init__.py
├── register.py
└── tasks
    ├── BUILD
    ├── __init__.py
    └── jupyter.py

The __init__.py are empty. I believe the BUILD files are not relevant.

jupyter/register.py:

from __future__ import (absolute_import, division, generators, nested_scopes,
                        print_function, unicode_literals, with_statement)

from chartbeat.pants.plugins.jupyter.tasks.jupyter import JupyterRepl
from pants.goal.task_registrar import TaskRegistrar as task


def register_goals():
      task(name='jupyter', action=JupyterRepl).install()

jupyter/tasks/jupyter.py:

from __future__ import (absolute_import, division, generators, nested_scopes, print_function,
                        unicode_literals, with_statement)

import os
import signal

from pex.pex_info import PexInfo

from pants.backend.python.python_requirement import PythonRequirement
from pants.backend.python.targets.python_requirement_library import PythonRequirementLibrary
from pants.backend.python.targets.python_target import PythonTarget
from pants.backend.python.tasks.python_execution_task_base import PythonExecutionTaskBase
from pants.task.repl_task_mixin import ReplTaskMixin
from pants.util.contextutil import signal_handler_as


class JupyterRepl(ReplTaskMixin, PythonExecutionTaskBase):
  """ Run a pants aware jupyter notebook """
  @classmethod
  def register_options(cls, register):
    super(JupyterRepl, cls).register_options(register)
    register('--jupyter-entry-point', advanced=True, default='jupyter_stubber.stubber:run',
             help='The Jupyter REPL entry point.')
    register('--jupyter-requirements', advanced=True, type=list, default=['jupyter_stubber==1.0.0'],
             help='Jupyter Requirements.')

  @classmethod
  def select_targets(cls, target):
    return isinstance(target, (PythonTarget, PythonRequirementLibrary))

  def extra_requirements(self):
    return self.get_options().jupyter_requirements

  def setup_repl_session(self, targets):
    entry_point = self.get_options().jupyter_entry_point
    pex_info = PexInfo.default()
    pex_info.entry_point = entry_point
    return self.create_pex(pex_info)

  # NB: **pex_run_kwargs is used by tests only.
  def launch_repl(self, pex, **pex_run_kwargs):
    def ignore_control_c(signum, frame): pass

    with signal_handler_as(signal.SIGINT, ignore_control_c):
      env = pex_run_kwargs.pop('env', os.environ).copy()
      pex.run(env=env, **pex_run_kwargs)

The jupyter_entry_point is some code for ensuring the process that runs jupyter kernel client inherits the correct import paths. Its BUILD file contains all the 3rdparty requirements for Jupyter.

Jupyter Stubber Code

Organization of the stubber code:

jupyter_stubber
├── BUILD
├── README.rst
├── __init__.py
└── stubber.py

The BUILD contents are as follows:

python_library(
    name = "jupyter_stubber",
    sources = ["stubber.py"],
    dependencies = [
        "3rdparty/python:ipykernel",
        "3rdparty/python:ipython",
        "3rdparty/python:jupyter-console",
        "3rdparty/python:nbconvert",
        "3rdparty/python:nbformat",
        "3rdparty/python:notebook",
        "3rdparty/python:pyzmq",
        "3rdparty/python:tornado",
    ],
    provides = setup_py(
        name='jupyter_stubber',
        version='1.0.0',
    ),
)

python_library(
    name = "jupyter_stubber_py3",
    sources = ["stubber.py"],
    dependencies = [
        "3rdparty/python3:ipykernel",
        "3rdparty/python3:ipython",
        "3rdparty/python3:jupyter-console",
        "3rdparty/python3:nbconvert",
        "3rdparty/python3:nbformat",
        "3rdparty/python3:notebook",
        "3rdparty/python3:pyzmq",
        "3rdparty/python3:tornado",
    ],
    provides = setup_py(
        name='jupyter_stubber_py3',
        version='1.0.0',
    ),
    compatibility="CPython>=3.6.8,<4",
)

The contents of stubber.py are as follows:

# -*- coding: utf-8 -*-
import os
import re
import sys

from notebook.notebookapp import main

def run():
    # Pants modifies sys.path to get all it's magic working. Jupyter's
    # kernel client uses Popen to spawn a kenerl in another process which
    # won't inherit the correct import paths unless we copy them into the
    # environment like we do here
    os.environ['PYTHONPATH'] = ":".join(sys.path)
    sys.argv[0] = re.sub(r'(-script\.pyw?|\.exe)?$', '', sys.argv[0])
    sys.exit(main())
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment