I have a context manager that captures output to a string for a block of code indented under a with
statement. This context manager yields a custom result object which will, when the block has finished executing, contain the captured output.
from contextlib import contextmanager
@contextmanager
def capturing():
"Captures output within a 'with' block."
from cStringIO import StringIO
class result(object):
def __init__(self):
self._result = None
def __str__(self):
return self._result
try:
stringio = StringIO()
out, err, sys.stdout, sys.stderr = sys.stdout, sys.stderr, stringio, stringio
output = result()
yield output
finally:
output._result, sys.stdout, sys.stderr = stringio.getvalue(), out, err
stringio.close()
with capturing() as text:
print "foo bar baz",
print str(text) # prints "foo bar baz"
I can't just return a string, of course, because strings are immutable and thus the one the user gets back from the with
statement can't be changed after their block of code runs. However, it is something of a drag to have to explicitly convert the result object to a string after the fact with str
(I also played with making the object callable as a bit of syntactic sugar).
So is it possible to make the result instance act like a string, in that it does in fact return a string when named? I tried implementing __get__
, but that appears to only work on attributes. Or is what I want to do not really possible?
At first glance, it looked like UserString
(well, actually MutableString
, but that's going away in Python 3.0) was basically what I wanted. Unfortunately, UserString doesn't work quite enough like a string; I was getting some odd formatting in print
statements ending in commas that worked fine with str
strings. (It appears you get an extra space printed if it's not a "real" string, or something.) I had the same issue with a toy class I created to play with wrapping a string. I didn't take the time to track down the cause, but it appears UserString
is most useful as an example.
I actually ended up using a bytearray
because it works enough like a string for most purposes, but is mutable. I also wrote a separate version that splitlines()
the text into a list. This works great and is actually better for my immediate use case, which is removing "extra" blank lines in the concatenated output of various functions. Here's that version:
import sys
from contextlib import contextmanager
@contextmanager
def capturinglines(output=None):
"Captures lines of output to a list."
from cStringIO import StringIO
try:
output = [] if output is None else output
stringio = StringIO()
out, err = sys.stdout, sys.stderr
sys.stdout, sys.stderr = stringio, stringio
yield output
finally:
sys.stdout, sys.stderr = out, err
output.extend(stringio.getvalue().splitlines())
stringio.close()
Usage:
with capturinglines() as output:
print "foo"
print "bar"
print output
['foo', 'bar']
with capturinglines(output): # append to existing list
print "baz"
print output
['foo', 'bar', 'baz']