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)
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())
%%autocmd
print(pwd())
pwd
Note that the tracebacks are pretty ugly, as is normal for IPython magics:
%%autocmd
print(ThisIsNotAProgram())
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)
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)
print(pandoc['-h']())
print(DontHaveThisProgramEither['-h']())
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.
No comments:
Post a Comment