Making IPython magics from click functions¶

Two fun facts:

click makes nice command-line APIs in Python
IPython uses magics to expose Python functions via a shell-like syntax

First, let's create the hello world click command from the click docs:

In [1]:

import click

@click.command()
@click.option('--count', default=1, help='Number of greetings.')
@click.option('--name', prompt='Your name',
              help='The person to greet.')
def hello(count, name):
    """Simple program that greets NAME for a total of COUNT times."""
    for x in range(count):
        click.echo('Hello %s!' % click.style(name, fg='green'))

which we can call:

In [2]:

hello(['--count', '2'])

Your name: min
Hello min!
Hello min!

An exception has occurred, use %tb to see the full traceback.

SystemExit: 0

/Users/benjaminrk/dev/ip/ipython/IPython/core/interactiveshell.py:3283: UserWarning: To exit: use 'exit', 'quit', or Ctrl-D.
  warn("To exit: use 'exit', 'quit', or Ctrl-D.", stacklevel=1)

Two things:

it raises SystemExit, and
min is not green

The latter is because click checks if sys.stdin is a tty, which it isn't in the notebook, and only enables colors if it is.

The check is in click.utils.should_strip_ansi:

In [3]:

s = click.style('Hello World!', fg='green')
s

Out[3]:

'\x1b[32mHello World!\x1b[0m'

In [4]:

click.utils.should_strip_ansi??

Signature: click.utils.should_strip_ansi(stream=None, color=None)
Docstring: <no docstring>
Source:   
def should_strip_ansi(stream=None, color=None):
    if color is None:
        if stream is None:
            stream = sys.stdin
        return not isatty(stream)
    return not color
File:      ~/conda/lib/python3.6/site-packages/click/_compat.py
Type:      function

In [5]:

click.echo(s)

Hello World!

We can force click not to check this and then echo should produce colored output:

In [6]:

try:
    from unittest import mock
except ImportError: # py2
    import mock

with mock.patch(
    'click.utils.should_strip_ansi',
    lambda *args, **kwargs: False
):
    click.echo(s)

Hello World!

Next, we can write the transform to turn a click command into an IPython magic. An IPython magic is a Python function called where the rest of the line is passed as a simple string. The click function wants a list of arguments, sys.argv-style.

So we need to do the following things:

split the line into arguments with shlex.split
tell the command that it's name is %magicname
tell click.echo that it should always color output
catch SystemExit errors and turning them into IPython magic UsageErrors
finally, register the magic via IPython's register_magic_function

In [7]:

import shlex
try:
    from unittest import mock
except ImportError: # py2
    import mock

from IPython.core.error import UsageError

def click_magic(click_command, name=None):
    if name is None:
        name = click_command.name

    def magic_func(line):
        args = shlex.split(line)
        # bypass click's check for whether colors should be enabled
        with mock.patch(
            'click.utils.should_strip_ansi',
            lambda *args, **kwargs: False
        ):
            try:
                click_command(
                    shlex.split(line),
                    prog_name='%' + name,
                )
            except SystemExit as e:
                if e.code != 0:
                    raise UsageError("Command exited with status=%s" % e.code)
    
    get_ipython().register_magic_function(magic_func, magic_name=name)

In [8]:

click_magic(hello)

In [9]:

%hello --help

Usage: %hello [OPTIONS]

  Simple program that greets NAME for a total of COUNT times.

Options:
  --count INTEGER  Number of greetings.
  --name TEXT      The person to greet.
  --help           Show this message and exit.

In [10]:

%hello --count 2 --name myname

Hello myname!
Hello myname!

In [11]:

%hello --unrecognized

Error: no such option: --unrecognized
UsageError: Command exited with status=2