How to Write a Jupyter Magic in Python

Jupyter magics allow us to run convenient utility functions within Jupyter notebooks. Anyone who has done much data analysis in a Jupyter notebook is likely familiar with

%matplotlib inline

which causes our matplotlib figures to be rendered in the notebook. This short post will explain the mechanics of creating Jupyter notebooks and exists mostly as a reference for my future self. For a slightly more involved example, my package giphy-ipython-magic serves well.

!pip install -q giphy-ipython-magic

%load_ext giphy_magic
%giphy magic

A simple magic

To start, we’ll implement a Jupyter magic that prints the result of cowsay (one of my favorite Unix utilities) given a phrase.

!pip install -q cowsay

%%writefile ./cowsay_magic.py
import cowsay as cs

def cowsay(msg):
    cs.cow(msg)

Overwriting ./cowsay_magic.py

Here the %%writefile magic writes the contents of the rest of the cell to the cowsay_magic.py file in the current directory. The script written to this file calls a Python library that reimplements cowsay and prints the result. In order for Jupyter to know that this file and function define a magic command, we must register the magic in a function named load_ipython_extension. (Note that we could also use the @register_line_magic decorator, but load_ipython_extension is necessary to redefine this magic momentarily. If anyone knows how to do this with the decorator, I’m all ears.)

%%writefile -a ./cowsay_magic.py
def load_ipython_extension(ipython):
    ipython.register_magic_function(cowsay, 'line')

Appending to ./cowsay_magic.py

Here the -a argument causes %%writefile to append to the existing file instead of overwriting it, which is the default behavior.

We make sure cowsay_magic.py is on the PYTHONPATH and load the magic into the Jupyter environment.

import sys

sys.path.append('.')

%load_ext cowsay_magic

We can now use %cowsay to summon our bovine friend.

%cowsay Hello Jupyter!

  ______________
| Hello Jupyter! |
  ==============
              \
               \
                 ^__^
                 (oo)\_______
                 (__)\       )\/\
                     ||----w |
                     ||     ||

Adding arguments

Jupyter passes the string after the magic as msg, and many magics implement shell-style arguments. We will add argument parsing to %cowsay in order to change the type of figure in the ASCII art.

%%writefile ./cowsay_magic.py
from argparse import ArgumentParser
import cowsay as cs

def parse_args(msg):
    parser = ArgumentParser(prog='cowsay magic')
    parser.add_argument('-f', dest='char_name', action='store', default='cow')
    parser.add_argument('message', nargs='*')
    
    return parser.parse_args(msg.split())

def cowsay(msg):
    args = parse_args(msg)
    
    print(cs.get_output_string(args.char_name, ' '.join(args.message)))
    
def load_ipython_extension(ipython):
    ipython.register_magic_function(cowsay, 'line')

Overwriting ./cowsay_magic.py

Here we have used the argparse module to parse msg. We reload the cowsay_magic extension.

%reload_ext cowsay_magic

Passing no arguments to %cowsay still prints a cow.

%cowsay Hello Jupyter!

  ______________
| Hello Jupyter! |
  ==============
              \
               \
                 ^__^
                 (oo)\_______
                 (__)\       )\/\
                     ||----w |
                     ||     ||

Passing the -f argument to %cowsay changes the speaker.

%cowsay -f trex Hello Jupyter!

  ______________
| Hello Jupyter! |
  ==============
                     \
                      \
                       \
                        \
                           .-=-==--==--.
                     ..-=="  ,'o`)      `.
                   ,'         `"'         \
                  :  (                     `.__...._
                  |                  )    /         `-=-.
                  :       ,vv.-._   /    /               `---==-._
                   \/\/\/VV ^ d88`;'    /                         `.
                       ``  ^/d88P!'    /             ,              `._
                          ^/    !'   ,.      ,      /                  "-,,__,,--'""""-.
                         ^/    !'  ,'  \ . .(      (         _           )  ) ) ) ))_,-.\
                        ^(__ ,!',"'   ;:+.:%:a.     \:.. . ,'          )  )  ) ) ,"'    '
                        ',,,'','     /o:::":%:%a.    \:.:.:         .    )  ) _,'
                         """'       ;':::'' `+%%%a._  \%:%|         ;.). _,-""
                                ,-='_.-'      ``:%::)  )%:|        /:._,"
                               (/(/"           ," ,'_,'%%%:       (_,'
                                              (  (//(`.___;        \
                                               \     \    `         `
                                                `.    `.   `.        :
                                                  \. . .\    : . . . :
                                                   \. . .:    `.. . .:
                                                    `..:.:\     \:...\
                                                     ;:.:.;      ::...:
                                                     ):%::       :::::;
                                                 __,::%:(        :::::
                                              ,;:%%%%%%%:        ;:%::
                                                ;,--""-.`\  ,=--':%:%:\
                                               /"       "| /-".:%%%%%%%\
                                                               ;,-"'`)%%)
                                                              /"      "|

Working with Python objects

Our %cowsay magic works only with strings, but we can also manipulate Python objects in a magic function using eval. To demonstrate, we will define a magic to invert the y-axis of a matplotlib plot.

%%writefile flip_magic.py
from IPython.core.magic import needs_local_scope

@needs_local_scope
def flip(fig_str, local_ns=None):
    fig = eval(fig_str, None, local_ns)
    fig.gca().invert_yaxis()
    
    return fig

def load_ipython_extension(ipython):
    ipython.register_magic_function(flip, 'line')

Overwriting flip_magic.py

Note the @needs_local_scope decorater that tells Jupyter to pass the local scope to our magic function. We load flip_magic and see that it does indeed invert the y-axis of a simple plot.

from matplotlib import pyplot as plt

fig, ax = plt.subplots(figsize=(8, 6))

ax.plot([0, 1], [0, 1]);

%load_ext flip_magic

%flip fig

I hope that this simple tutorial has been helpful. For more detail about the custom magic API, consult the excellent Jupyter documentation.

The notebook this post was generated from is available as a Jupyter notebook here.