Capturing output of print in Python

The standard Python disassembler, dis, prints its output directly to standard output instead of returning it as a string. This is very inconvenient if you wish to do further processing on the disassembled code, or even if you just want to do something simple like count the number of lines.

Fortunately, it is easy to write a wrapper to capture the output of any Python function:

import sys, cStringIO, traceback

def capture(func, *args, **kwargs):
    """Capture the output of func when called with the given arguments.

    The function output includes any exception raised. capture returns
    a tuple of (function result, standard output, standard error).
    """
    stdout, stderr = sys.stdout, sys.stderr
    sys.stdout = c1 = cStringIO.StringIO()
    sys.stderr = c2 = cStringIO.StringIO()
    result = None
    try:
        result = func(*args, **kwargs)
    except:
        traceback.print_exc()
    sys.stdout = stdout
    sys.stderr = stderr
    return (result, c1.getvalue(), c2.getvalue())


With the aid of capture it is easy to grab the disassembled code:

import dis

def disassemble(obj=None):
    """Capture the output of dis.dis and return it."""
    return capture(dis.dis, obj)[1]