1from cStringIO import StringIO 2 3class StringIOTree(object): 4 """ 5 See module docs. 6 """ 7 8 def __init__(self, stream=None): 9 self.prepended_children = [] 10 if stream is None: 11 stream = StringIO() 12 self.stream = stream 13 self.write = stream.write 14 self.markers = [] 15 16 def getvalue(self): 17 content = [x.getvalue() for x in self.prepended_children] 18 content.append(self.stream.getvalue()) 19 return "".join(content) 20 21 def copyto(self, target): 22 """Potentially cheaper than getvalue as no string concatenation 23 needs to happen.""" 24 for child in self.prepended_children: 25 child.copyto(target) 26 stream_content = self.stream.getvalue() 27 if stream_content: 28 target.write(stream_content) 29 30 def commit(self): 31 # Save what we have written until now so that the buffer 32 # itself is empty -- this makes it ready for insertion 33 if self.stream.tell(): 34 self.prepended_children.append(StringIOTree(self.stream)) 35 self.prepended_children[-1].markers = self.markers 36 self.markers = [] 37 self.stream = StringIO() 38 self.write = self.stream.write 39 40 def insert(self, iotree): 41 """ 42 Insert a StringIOTree (and all of its contents) at this location. 43 Further writing to self appears after what is inserted. 44 """ 45 self.commit() 46 self.prepended_children.append(iotree) 47 48 def insertion_point(self): 49 """ 50 Returns a new StringIOTree, which is left behind at the current position 51 (it what is written to the result will appear right before whatever is 52 next written to self). 53 54 Calling getvalue() or copyto() on the result will only return the 55 contents written to it. 56 """ 57 # Save what we have written until now 58 # This is so that getvalue on the result doesn't include it. 59 self.commit() 60 # Construct the new forked object to return 61 other = StringIOTree() 62 self.prepended_children.append(other) 63 return other 64 65 def allmarkers(self): 66 children = self.prepended_children 67 return [m for c in children for m in c.allmarkers()] + self.markers 68 69 70__doc__ = r""" 71Implements a buffer with insertion points. When you know you need to 72"get back" to a place and write more later, simply call insertion_point() 73at that spot and get a new StringIOTree object that is "left behind". 74 75EXAMPLE: 76 77>>> a = StringIOTree() 78>>> a.write('first\n') 79>>> b = a.insertion_point() 80>>> a.write('third\n') 81>>> b.write('second\n') 82>>> a.getvalue().split() 83['first', 'second', 'third'] 84 85>>> c = b.insertion_point() 86>>> d = c.insertion_point() 87>>> d.write('alpha\n') 88>>> b.write('gamma\n') 89>>> c.write('beta\n') 90>>> b.getvalue().split() 91['second', 'alpha', 'beta', 'gamma'] 92>>> i = StringIOTree() 93>>> d.insert(i) 94>>> i.write('inserted\n') 95>>> out = StringIO() 96>>> a.copyto(out) 97>>> out.getvalue().split() 98['first', 'second', 'alpha', 'inserted', 'beta', 'gamma', 'third'] 99""" 100