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