Tuesday, April 12, 2011

How to Capture stdout into a Variable

To idea to capture stdout into a variable is relatively straightforward. What we need is basically to save the existing stdout into a temporary variable and then replace the existing stdout with another kind of stdout and then put the original stdout back. Below is how to achieve that in both Python and Java.

For Python:
def capture(func, *args, **kwargs):
    import sys, cStringIO
    
    tmpstdout = sys.stdout
    sys.stdout = cStringIO.StringIO()
    try:
        func(*args, **kwargs)
    finally:
        value = sys.stdout.getvalue()
        sys.stdout = tmpstdout
    return value

To use it:
def sayHello(): print "Hello World"

value = capture(sayHello)
print value.upper()

The output is:
HELLO WORLD

For Java:
package myproject;

import java.io.ByteArrayOutputStream;
import java.io.PrintStream;

public abstract class StdoutCaptor {
    
    public String capture() {
        PrintStream tmp = System.out;
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        System.setOut(new PrintStream(baos));
        try {
            invoke();
        } finally {
            System.setOut(tmp);
        }
        return new String(baos.toByteArray());
    }
    
    public abstract void invoke();
}

To use it:
package myproject;

public class Main {
    
    public static void sayHello() {
        System.out.println("Hello World!");
    }
    
    public static void main(String[] args) {
        String value = new StdoutCaptor() {
            @Override
            public void invoke() {
                sayHello();
            }
        }.capture();
        System.out.println(value.toUpperCase());
    }
}

The output is:
HELLO WORLD

Unlike Python, Java doesn't allow a function/method to be passed around. Hence, the code is more verbose than Python.

No comments:

Post a Comment