Another Way to Capture Print in Python

When working with the Python ast module i found the pycodegen which generates bytecode from an ast. Basically what you can do with it is modifying the ast before compiling the code. So one use case is compiling code and patching all Print/Printnl nodes without a destination to buffer into a stream:

import compiler

def compile_with_stream(code, stream, filename='?'):
    print_nodes = [compiler.ast.Print, compiler.ast.Printnl]
    rv = compiler.parse(code)
    nodes = [rv]
    while nodes:
        node = nodes.pop()
        node.filename = filename
        if node.__class__ in print_nodes and node.dest is None:
            node.dest = compiler.ast.Const(stream, node.lineno)
        nodes.extend(node.getChildNodes())
    gen = compiler.pycodegen.ModuleCodeGenerator(rv)
    return gen.getCode()

Instead of using compiler.ast.Const you can also use compiler.ast.Name and point to a variable that contains the stream. This would allow caching the bytecode too. Now what it does is translating “print ‘Hello World’” to “print >> stream, ‘Hello World’” if there is no given destination. This can be useful for template engines (once again ^^).

Example usage:

from StringIO import StringIO
stream = StringIO()
exec compile_with_stream('print "Hello World"', stream)
print 'Captured: %r' % stream.getvalue()

Leave a Reply

cogitations driven by wordpress