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!")
>>> from plumbum import local >>> echo = local['echo'] >>> echo("This is echoed!")
>>> 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:
Note that the tracebacks are pretty ugly, as is normal for IPython magics:
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:
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.