14
Jul
09

hacking python at runtime:a cool way of modifying your scripts



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[0] will be the import statement, and ast.body[1] 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[0])
os_import.names[0].name = "os"


This object is now equivalent to the following code:


import os


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[0])
# 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[1]
# change the argument of listdir to D:\
for_obj.iter.args[0].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 :


exec code

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[0])
	# remove the import geo statement
	ast.body.remove(ast.body[0])
	# change the import argument to os
	os_import.names[0].name = "os"
	# add the import os statement to the ast
	ast.body.insert(0,os_import)
	for_obj = ast.body[1]
	# change the argument of listdir to D:\
	for_obj.iter.args[0].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.

About these ads

4 Responses to “hacking python at runtime:a cool way of modifying your scripts”


  1. July 15, 2009 at 19:38

    This is 2.6 and up trick only :(. Nice work :)

  2. July 19, 2009 at 10:05

    Good to see Python showing more of its Lisp heritage, hope it gets to be more syntactically obvious like macros are in Lisp.

  3. 3 Bob/Paul
    November 18, 2010 at 20:11

    I’m failing to see how this is really useful.

    I mean, yes, it’s quite cool. But, you’re less modifying the script at runtime and more modifying it immediately before runtime. Perhaps it’s just I’m stuck on the example you provided, but why wouldn’t you just update the script as it exists on disk so it imports os instead of geo? What is the benefit of leaving the original script broken and fixing it with a special wrapper designed specifically for fixing that exact code? I mean, I suppose you could loop through the code searching ast.body[x].names[y].name == ‘geo’ and ast.body[x] is an import object, but you’d still have to know that error was coming. Why not just edit the original script before you run it?


Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s


Blog Stats

  • 176,398 hits

Follow

Get every new post delivered to your Inbox.