#!/usr/bin/env python
# coding: utf-8

# ---
# title: How to Write a Jupyter Magic in Python
# tags: Python, Jupyter
# ---

# [Jupyter magics](https://ipython.readthedocs.io/en/stable/interactive/magics.html) 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

# In[1]:


get_ipython().run_line_magic('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`](https://github.com/AustinRochford/giphy-ipython-magic) serves well.

# In[2]:


get_ipython().system('pip install -q giphy-ipython-magic')


# In[3]:


get_ipython().run_line_magic('load_ext', 'giphy_magic')
get_ipython().run_line_magic('giphy', 'magic')


# ## A simple magic
# 
# To start, we'll implement a Jupyter magic that prints the result of [`cowsay`](https://en.wikipedia.org/wiki/Cowsay) (one of my favorite Unix utilities) given a phrase.

# In[4]:


get_ipython().system('pip install -q cowsay')


# In[5]:


get_ipython().run_cell_magic('writefile', './cowsay_magic.py', 'import cowsay as cs\n\ndef cowsay(msg):\n    cs.cow(msg)\n')


# Here the [`%%writefile` magic](https://ipython.readthedocs.io/en/stable/interactive/magics.html#cellmagic-writefile) 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](https://github.com/VaasuDevanS/cowsay-python) 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](http://localhost:8888/notebooks/jupyter_magic/Jupyter%20Magic.ipynb), 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.)

# In[6]:


get_ipython().run_cell_magic('writefile', '-a ./cowsay_magic.py', "def load_ipython_extension(ipython):\n    ipython.register_magic_function(cowsay, 'line')\n")


# 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.

# In[7]:


import sys

sys.path.append('.')


# In[8]:


get_ipython().run_line_magic('load_ext', 'cowsay_magic')


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

# In[9]:


get_ipython().run_line_magic('cowsay', 'Hello Jupyter!')


# ### 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.

# In[10]:


get_ipython().run_cell_magic('writefile', './cowsay_magic.py', "from argparse import ArgumentParser\nimport cowsay as cs\n\ndef parse_args(msg):\n    parser = ArgumentParser(prog='cowsay magic')\n    parser.add_argument('-f', dest='char_name', action='store', default='cow')\n    parser.add_argument('message', nargs='*')\n    \n    return parser.parse_args(msg.split())\n\ndef cowsay(msg):\n    args = parse_args(msg)\n    \n    print(cs.get_output_string(args.char_name, ' '.join(args.message)))\n    \ndef load_ipython_extension(ipython):\n    ipython.register_magic_function(cowsay, 'line')\n")


# Here we have used the [`argparse`](https://docs.python.org/3/library/argparse.html) module to parse `msg`.  We reload the `cowsay_magic` extension.

# In[11]:


get_ipython().run_line_magic('reload_ext', 'cowsay_magic')


# Passing no arguments to `%cowsay` still prints a cow.

# In[12]:


get_ipython().run_line_magic('cowsay', 'Hello Jupyter!')


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

# In[13]:


get_ipython().run_line_magic('cowsay', '-f trex Hello Jupyter!')


# ## Working with Python objects
# 
# Our `%cowsay` magic works only with strings, but we can also manipulate Python objects in a magic function using [`eval`](https://docs.python.org/3/library/functions.html#eval).  To demonstrate, we will define a magic to invert the y-axis of a `matplotlib` plot.

# In[14]:


get_ipython().run_cell_magic('writefile', 'flip_magic.py', "from IPython.core.magic import needs_local_scope\n\n@needs_local_scope\ndef flip(fig_str, local_ns=None):\n    fig = eval(fig_str, None, local_ns)\n    fig.gca().invert_yaxis()\n    \n    return fig\n\ndef load_ipython_extension(ipython):\n    ipython.register_magic_function(flip, 'line')\n")


# 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.

# In[15]:


from matplotlib import pyplot as plt


# In[16]:


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

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


# In[17]:


get_ipython().run_line_magic('load_ext', 'flip_magic')


# In[18]:


get_ipython().run_line_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](https://ipython.readthedocs.io/en/stable/config/custommagics.html).
# 
# The notebook this post was generated from is available as a Jupyter notebook [here](https://nbviewer.jupyter.org/gist/AustinRochford/6c48e60953fad4069ff027c3fcdccc9a).