"""
OutputBuffer: A port of some of PHP's ob_* functions.
sys.stdout is redirected to StringIO buffers.  Buffers
can be stacked, and buffers can optionally be pop()'ed
without flushing their contents.

Usage:

from OutputBuffer import OutputBuffer
ob = OutputBuffer() #sys.stdout unchanged yet
ob.push()           #Create a new buffer
print "abc"         #Output redirected to topmost buffer
ob.push()           #Push another buffer
print "def"         #Output redirected to topmost buffer
ob.pop()            #Pop a buffer, flush first
s = str(ob)         #Current value of topmost buffer
ob.pop(flush=0)     #Pop a buffer, don't flush first
ob.pop()            #raise IOError when pop too many times

When the first push() is pop()'ed, sys.stdout will
be set to its original value.

"""

__author__ = "Patrick Lioi <patrick_at_lioi_dot_net>"
__date__ = "2003/09/16 4:16:24"
__version__ = "1.0"

from StringIO import StringIO
import sys

class OutputBuffer:
    def __init__(self):
        #Stack always contains at least sys.stdout
        self.streams = [sys.stdout]
    def push(self):
        self.streams.append(StringIO())
        sys.stdout = self.streams[-1]
    def __str__(self):
        return self.streams[-1].getvalue()
    def pop(self, flush=1):
        if (len(self.streams)==1):
            raise IOError, "pop() called on empty buffer stack"
        if flush:
            s = self.__str__()
        self.streams[-1].close()
        self.streams.pop()
        sys.stdout = self.streams[-1]
        if flush:
            print s, #Suppress extra endline

if __name__ == "__main__":
    ob = OutputBuffer()
    print "Printing to regular sys.stdout."
    ob.push()
    print "Printing to first buffer."
    ob.push()
    print "Printing to second buffer."
    ob.push()
    print "Printing to third buffer.  This buffer will not be output upon pop()."
    s = str(ob)     #store topmost before pop(flush=0)
    ob.pop(flush=0) #pop third buffer
    ob.pop()        #pop second buffer
    ob.pop()        #pop first buffer
    print "Regular sys.stdout has been restored."
    ob.pop()        #raise IOError