Making an autoload extension for IPython

I recently decided to try my hand at making an auto-load extension for Python and Plumbum. I was planning to suggest it as a new feature, then I thought it might be an experimental feature, and now it’s just a blog post. But it was an interesting idea and didn’t seem to be well documented process on the web. So, here it is.

The plan was to make commands like this:

>>> from plumbum.cmd import echo
>>> echo("This is echoed!")

or

>>> from plumbum import local
>>> echo = local['echo']
>>> echo("This is echoed!")

into this:

>>> echo("This is echoed!")

Thereby making Plumbum even more like Bash for fast scripting.

First attempt: Magic

I originally thought it would be a magic:

%%writefile local_magic.py
try:
    from IPython.core.magic import (Magics, magics_class,
                                cell_magic, needs_local_scope)
except ImportError:
    print("IPython required for the IPython extension to be loaded.")
    raise

import ast
from plumbum import local, CommandNotFound
try:
    import builtins
except ImportError: # Python 2
    import __builtins__ as builtins

@magics_class
class AutoMagics(Magics):

    @needs_local_scope
    @cell_magic
    def autocmd(self, line, cell, local_ns=None):
        mod = ast.parse(cell)
        calls = [c for c in ast.walk(mod) if isinstance(c,ast.Call) or isinstance(c, ast.Subscript)]
        for call in calls:
            name = call.func.id if isinstance(call, ast.Call) else call.value.id
            if name not in self.shell.user_ns and name not in dir(builtins):
                try:
                    self.shell.user_ns[name] = local[name]
                except CommandNotFound:
                    pass
        exec(cell, self.shell.user_ns, local_ns)

def load_ipython_extension(ipython):
    ipython.register_magics(AutoMagics)
Writing local_magic.py

Here, I grab the contents of a cell, which is still a string, and run it through the python AST. I look at the nodes produced, and for every Call or Subscript node, I check and see if it’s in the current namespace. If it’s not, I make one attepmt to load it from plumbum.local, and pass if I can’t. Then, I exec the cell as normal.

Now, let’s see if it works:

%load_ext local_magic
print(pwd())
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
 in ()
----> 1 print(pwd())

NameError: name 'pwd' is not defined
%%autocmd
print(pwd())
/home/henryiii/Dropbox/Personal/notebooks
pwd
LocalCommand()

Note that the tracebacks are pretty ugly, as is normal for IPython magics:

%%autocmd
print(ThisIsNotAProgram())
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
 in ()
----> 1 get_ipython().run_cell_magic('autocmd', '', 'print(ThisIsNotAProgram())')

/usr/local/lib/python3.4/dist-packages/IPython/core/interactiveshell.py in run_cell_magic(self, magic_name, line, cell)
   2262             magic_arg_s = self.var_expand(line, stack_depth)
   2263             with self.builtin_trap:
-> 2264                 result = fn(magic_arg_s, cell)
   2265             return result
   2266 

/home/henryiii/Dropbox/Personal/notebooks/local_magic.py in autocmd(self, line, cell, local_ns)

/usr/local/lib/python3.4/dist-packages/IPython/core/magic.py in (f, *a, **k)
    191     # but it's overkill for just that one bit of state.
    192     def magic_deco(arg):
--> 193         call = lambda f, *a, **k: f(*a, **k)
    194 
    195         if callable(arg):

/home/henryiii/Dropbox/Personal/notebooks/local_magic.py in autocmd(self, line, cell, local_ns)
     28                 except CommandNotFound:
     29                     pass
---> 30         exec(cell, self.shell.user_ns, local_ns)
     31 
     32 def load_ipython_extension(ipython):

 in ()

NameError: name 'ThisIsNotAProgram' is not defined

Better Than Magic

We didn’t really gain much, though. We still have to type the extra line, we just don’t have to repeat the name of the function we are importing.

How about removing the silly requirement that we stick a magic in front of every command? Easy, we can use a hook in IPython to get the AST transform:

%%writefile local_ext.py
import IPython, ast, builtins
from plumbum import local, CommandNotFound, FG, BG, TF, RETCODE

class NameLoader(ast.NodeTransformer):
    """Direct calls to functions not defined are imported from local if possible."""
    def visit_Call(self, node):
        ipython = IPython.get_ipython()
        if isinstance(node.func, ast.Name): # Will ignore names with .
            name = node.func.id
            if name not in ipython.user_ns and name not in dir(builtins):
                try:
                    ipython.user_ns[name] = local[name]
                except CommandNotFound:
                    pass
        self.generic_visit(node) # The rest of the nodes should be visited
        return node
    
    def visit_Subscript(self, node):
        """Bracket syntax too."""
        ipython = IPython.get_ipython()
        if isinstance(node.value, ast.Name): # Will ignore names with .
            name = node.value.id
            if name not in ipython.user_ns and name not in dir(builtins):
                try:
                    ipython.user_ns[name] = local[name]
                except CommandNotFound:
                    pass
        self.generic_visit(node) # The rest of the nodes should be visited
        return node
    
ipython = IPython.get_ipython()
nl = NameLoader()
    
def load_ipython_extension(ipython):
    ipython.ast_transformers.append(nl)
    ipython.user_ns['local'] = local
    ipython.user_ns['BG'] = BG
    ipython.user_ns['FG'] = FG
    ipython.user_ns['TF'] = TF
    ipython.user_ns['RETCODE'] = RETCODE
    
def unload_ipython_extension(ipython):
    ipython.ast_transformers.remove(nl)
Overwriting local_ext.py

Here, ast_transformers will allow us to run on the AST tree after IPython has converted IPython syntax out (which is nicer than before, as IPython code still works), and it a bit cleaner. All tracebacks should be normal now. We also go ahead an populate the namespace of IPython, so local and some modifiers are available.

We load it, causing all commands that are not found to be loaded from local if possible:

%load_ext local_ext
echo(2)
'2\n'
print(pandoc['-h']())
pandoc [OPTIONS] [FILES]
Input formats:  docbook, haddock, html, json, latex, markdown, markdown_github,
                markdown_mmd, markdown_phpextra, markdown_strict, mediawiki,
                native, opml, org, rst, textile
Output formats: asciidoc, beamer, context, docbook, docx, dzslides, epub, epub3,
                fb2, html, html5, icml, json, latex, man, markdown,
                markdown_github, markdown_mmd, markdown_phpextra,
                markdown_strict, mediawiki, native, odt, opendocument, opml,
                org, pdf*, plain, revealjs, rst, rtf, s5, slideous, slidy,
                texinfo, textile
                [*for pdf output, use latex or beamer and -o FILENAME.pdf]
Options:
  -f FORMAT, -r FORMAT  --from=FORMAT, --read=FORMAT                    
  -t FORMAT, -w FORMAT  --to=FORMAT, --write=FORMAT                     
  -o FILENAME           --output=FILENAME                               
                        --data-dir=DIRECTORY                            
                        --strict                                        
  -R                    --parse-raw                                     
  -S                    --smart                                         
                        --old-dashes                                    
                        --base-header-level=NUMBER                      
                        --indented-code-classes=STRING                  
  -F PROGRAM            --filter=PROGRAM                                
                        --normalize                                     
  -p                    --preserve-tabs                                 
                        --tab-stop=NUMBER                               
  -s                    --standalone                                    
                        --template=FILENAME                             
  -M KEY[:VALUE]        --metadata=KEY[:VALUE]                          
  -V KEY[:VALUE]        --variable=KEY[:VALUE]                          
  -D FORMAT             --print-default-template=FORMAT                 
                        --print-default-data-file=FILE                  
                        --no-wrap                                       
                        --columns=NUMBER                                
                        --toc, --table-of-contents                      
                        --toc-depth=NUMBER                              
                        --no-highlight                                  
                        --highlight-style=STYLE                         
  -H FILENAME           --include-in-header=FILENAME                    
  -B FILENAME           --include-before-body=FILENAME                  
  -A FILENAME           --include-after-body=FILENAME                   
                        --self-contained                                
                        --offline                                       
  -5                    --html5                                         
                        --html-q-tags                                   
                        --ascii                                         
                        --reference-links                               
                        --atx-headers                                   
                        --chapters                                      
  -N                    --number-sections                               
                        --number-offset=NUMBERS                         
                        --no-tex-ligatures                              
                        --listings                                      
  -i                    --incremental                                   
                        --slide-level=NUMBER                            
                        --section-divs                                  
                        --default-image-extension=extension             
                        --email-obfuscation=none|javascript|references  
                        --id-prefix=STRING                              
  -T STRING             --title-prefix=STRING                           
  -c URL                --css=URL                                       
                        --reference-odt=FILENAME                        
                        --reference-docx=FILENAME                       
                        --epub-stylesheet=FILENAME                      
                        --epub-cover-image=FILENAME                     
                        --epub-metadata=FILENAME                        
                        --epub-embed-font=FILE                          
                        --epub-chapter-level=NUMBER                     
                        --latex-engine=PROGRAM                          
                        --bibliography=FILE                             
                        --csl=FILE                                      
                        --citation-abbreviations=FILE                   
                        --natbib                                        
                        --biblatex                                      
  -m[URL]               --latexmathml[=URL], --asciimathml[=URL]        
                        --mathml[=URL]                                  
                        --mimetex[=URL]                                 
                        --webtex[=URL]                                  
                        --jsmath[=URL]                                  
                        --mathjax[=URL]                                 
                        --gladtex                                       
                        --trace                                         
                        --dump-args                                     
                        --ignore-args                                   
  -v                    --version                                       
  -h                    --help                                          
print(DontHaveThisProgramEither['-h']())
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
 in ()
----> 1 print(DontHaveThisProgramEither['-h']())

NameError: name 'DontHaveThisProgramEither' is not defined

This allows us to write code very quickly in a shell like environment by simply putting:

#!/usr/bin/env ipython
%load_ext local_ext

at the top of our scripts. (Assuming this is in the IPython path) Once the script is done, you can go import the programs used manually and then remove the IPython parts.

comments powered by Disqus