Search code examples
pythonstringwith-statement

How to make a class that acts like a string?


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?


Solution

  • 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']