I recently discovered the _ast module. Using it, one can process Python’s syntax trees.
I’m going to illustrate the use of this module with a simple example. Here’s the code snippet we’re going to work on :
import geo for file in os.listdir("."): print file
Let’s talk a bit about this snippet.
It’s kind of obvious this script has runtime errors. On my Python installation there’s no module named geo … chances are there’s not one on your system either 🙂
Another error would be the call to os.listdir. This is not an error per-se, because if you import the os module, no exception will get thrown at runtime. We can easily notice that the os module has not been imported.
Here’s what we’ll accomplish using the _ast module :
- remove the import of the module geo
- add an import to the os module
- to make things more interesting, we’ll change the “.” argument of the listdir function to the “D:\\” drive
So, how do we get the syntax tree? Pretty simple :
import _ast source_code = """ import geo for file in os.listdir("."): print file """ ast = compile(source_code,"<string>","exec",_ast.PyCF_ONLY_AST)
The compile function will build an AST object. The members of the AST can be accessed through the ast.body list. In this example, ast.body will be the import statement, and ast.body will be the for statement.
So, now we have an AST!
The first thing we’ll do is clone the import object. I’m doing this so that I don’t have to create an import object manually. I’m lazy, I know. If you don’t know how to clone a Python object, the following snippet illustrates it :
import copy an_object_copy = copy.deepcopy(an_object)
With this import clone, I want to import the os module. But, since this clone still has geo as it’s argument, we need to change that. We change that with the following snippet :
os_import = copy.deepcopy(ast.body) os_import.names.name = "os"
This object is now equivalent to the following code:
This is nice, but we have to add it to the AST. I’ll take advantage of this to remove the import geo statement:
# remove the import geo statement ast.body.remove(ast.body) # we insert the import os as the first statement ast.body.insert(0,os_import)
Right now, the code is runnable. No exceptions will be thrown. Before we run it, let’s accomplish the final task too. Let’s modify the argument of listdir from “.” to “D:\\”.We know that the for object will be the second object in the list:
for_obj = ast.body # change the argument of listdir to D:\ for_obj.iter.args.s = "D:\\"
This changes the argument. In case the attributes I’m setting seem magic, you can find them out using Python’s introspection system ( it’s how I found them too ). You can use this system from ipython, or even your python interpreter, by calling the dir function on any object. This call will list the name of the methods the object has.
Now that we have modified the AST, we need to transform it into runnable code. We do that by calling the compile function :
code = compile(ast,"<string>","exec")
We can run the code with the exec function :
Here’s the full code of the script:
import _ast import copy def fix_source(source_string): ast = compile(source_string,"<string>","exec",_ast.PyCF_ONLY_AST) # clone the import object, so we can modify it os_import = copy.deepcopy(ast.body) # remove the import geo statement ast.body.remove(ast.body) # change the import argument to os os_import.names.name = "os" # add the import os statement to the ast ast.body.insert(0,os_import) for_obj = ast.body # change the argument of listdir to D:\ for_obj.iter.args.s = "D:\\" # transform the AST into something runnable return compile(ast,"<string>","exec") if __name__ == "__main__": source_code = """ import geo for file in os.listdir("."): print file """ code = fix_source(source_code) exec code
I’m sure you can put this trick to use. It’s one of the greatest “hacks” I know.
I’ll try to post a Ruby alternative as soon as time allows me.