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