1r"""File-like objects that read from or write to a string buffer.
2
3This implements (nearly) all stdio methods.
4
5f = StringIO()      # ready for writing
6f = StringIO(buf)   # ready for reading
7f.close()           # explicitly release resources held
8flag = f.isatty()   # always false
9pos = f.tell()      # get current position
10f.seek(pos)         # set current position
11f.seek(pos, mode)   # mode 0: absolute; 1: relative; 2: relative to EOF
12buf = f.read()      # read until EOF
13buf = f.read(n)     # read up to n bytes
14buf = f.readline()  # read until end of line ('\n') or EOF
15list = f.readlines()# list of f.readline() results until EOF
16f.truncate([size])  # truncate file at to at most size (default: current pos)
17f.write(buf)        # write at current position
18f.writelines(list)  # for line in list: f.write(line)
19f.getvalue()        # return whole file's contents as a string
20
21Notes:
22- Using a real file is often faster (but less convenient).
23- There's also a much faster implementation in C, called cStringIO, but
24  it's not subclassable.
25- fileno() is left unimplemented so that code which uses it triggers
26  an exception early.
27- Seeking far beyond EOF and then writing will insert real null
28  bytes that occupy space in the buffer.
29- There's a simple test set (see end of this file).
30"""
31try:
32    from errno import EINVAL
33except ImportError:
34    EINVAL = 22
35
36__all__ = ["StringIO"]
37
38def _complain_ifclosed(closed):
39    if closed:
40        raise ValueError, "I/O operation on closed file"
41
42class StringIO:
43    """class StringIO([buffer])
44
45    When a StringIO object is created, it can be initialized to an existing
46    string by passing the string to the constructor. If no string is given,
47    the StringIO will start empty.
48
49    The StringIO object can accept either Unicode or 8-bit strings, but
50    mixing the two may take some care. If both are used, 8-bit strings that
51    cannot be interpreted as 7-bit ASCII (that use the 8th bit) will cause
52    a UnicodeError to be raised when getvalue() is called.
53    """
54    def __init__(self, buf = ''):
55        # Force self.buf to be a string or unicode
56        if not isinstance(buf, basestring):
57            buf = str(buf)
58        self.buf = buf
59        self.len = len(buf)
60        self.buflist = []
61        self.pos = 0
62        self.closed = False
63        self.softspace = 0
64
65    def __iter__(self):
66        return self
67
68    def next(self):
69        """A file object is its own iterator, for example iter(f) returns f
70        (unless f is closed). When a file is used as an iterator, typically
71        in a for loop (for example, for line in f: print line), the next()
72        method is called repeatedly. This method returns the next input line,
73        or raises StopIteration when EOF is hit.
74        """
75        _complain_ifclosed(self.closed)
76        r = self.readline()
77        if not r:
78            raise StopIteration
79        return r
80
81    def close(self):
82        """Free the memory buffer.
83        """
84        if not self.closed:
85            self.closed = True
86            del self.buf, self.pos
87
88    def isatty(self):
89        """Returns False because StringIO objects are not connected to a
90        tty-like device.
91        """
92        _complain_ifclosed(self.closed)
93        return False
94
95    def seek(self, pos, mode = 0):
96        """Set the file's current position.
97
98        The mode argument is optional and defaults to 0 (absolute file
99        positioning); other values are 1 (seek relative to the current
100        position) and 2 (seek relative to the file's end).
101
102        There is no return value.
103        """
104        _complain_ifclosed(self.closed)
105        if self.buflist:
106            self.buf += ''.join(self.buflist)
107            self.buflist = []
108        if mode == 1:
109            pos += self.pos
110        elif mode == 2:
111            pos += self.len
112        self.pos = max(0, pos)
113
114    def tell(self):
115        """Return the file's current position."""
116        _complain_ifclosed(self.closed)
117        return self.pos
118
119    def read(self, n = -1):
120        """Read at most size bytes from the file
121        (less if the read hits EOF before obtaining size bytes).
122
123        If the size argument is negative or omitted, read all data until EOF
124        is reached. The bytes are returned as a string object. An empty
125        string is returned when EOF is encountered immediately.
126        """
127        _complain_ifclosed(self.closed)
128        if self.buflist:
129            self.buf += ''.join(self.buflist)
130            self.buflist = []
131        if n is None or n < 0:
132            newpos = self.len
133        else:
134            newpos = min(self.pos+n, self.len)
135        r = self.buf[self.pos:newpos]
136        self.pos = newpos
137        return r
138
139    def readline(self, length=None):
140        r"""Read one entire line from the file.
141
142        A trailing newline character is kept in the string (but may be absent
143        when a file ends with an incomplete line). If the size argument is
144        present and non-negative, it is a maximum byte count (including the
145        trailing newline) and an incomplete line may be returned.
146
147        An empty string is returned only when EOF is encountered immediately.
148
149        Note: Unlike stdio's fgets(), the returned string contains null
150        characters ('\0') if they occurred in the input.
151        """
152        _complain_ifclosed(self.closed)
153        if self.buflist:
154            self.buf += ''.join(self.buflist)
155            self.buflist = []
156        i = self.buf.find('\n', self.pos)
157        if i < 0:
158            newpos = self.len
159        else:
160            newpos = i+1
161        if length is not None and length >= 0:
162            if self.pos + length < newpos:
163                newpos = self.pos + length
164        r = self.buf[self.pos:newpos]
165        self.pos = newpos
166        return r
167
168    def readlines(self, sizehint = 0):
169        """Read until EOF using readline() and return a list containing the
170        lines thus read.
171
172        If the optional sizehint argument is present, instead of reading up
173        to EOF, whole lines totalling approximately sizehint bytes (or more
174        to accommodate a final whole line).
175        """
176        total = 0
177        lines = []
178        line = self.readline()
179        while line:
180            lines.append(line)
181            total += len(line)
182            if 0 < sizehint <= total:
183                break
184            line = self.readline()
185        return lines
186
187    def truncate(self, size=None):
188        """Truncate the file's size.
189
190        If the optional size argument is present, the file is truncated to
191        (at most) that size. The size defaults to the current position.
192        The current file position is not changed unless the position
193        is beyond the new file size.
194
195        If the specified size exceeds the file's current size, the
196        file remains unchanged.
197        """
198        _complain_ifclosed(self.closed)
199        if size is None:
200            size = self.pos
201        elif size < 0:
202            raise IOError(EINVAL, "Negative size not allowed")
203        elif size < self.pos:
204            self.pos = size
205        self.buf = self.getvalue()[:size]
206        self.len = size
207
208    def write(self, s):
209        """Write a string to the file.
210
211        There is no return value.
212        """
213        _complain_ifclosed(self.closed)
214        if not s: return
215        # Force s to be a string or unicode
216        if not isinstance(s, basestring):
217            s = str(s)
218        spos = self.pos
219        slen = self.len
220        if spos == slen:
221            self.buflist.append(s)
222            self.len = self.pos = spos + len(s)
223            return
224        if spos > slen:
225            self.buflist.append('\0'*(spos - slen))
226            slen = spos
227        newpos = spos + len(s)
228        if spos < slen:
229            if self.buflist:
230                self.buf += ''.join(self.buflist)
231            self.buflist = [self.buf[:spos], s, self.buf[newpos:]]
232            self.buf = ''
233            if newpos > slen:
234                slen = newpos
235        else:
236            self.buflist.append(s)
237            slen = newpos
238        self.len = slen
239        self.pos = newpos
240
241    def writelines(self, iterable):
242        """Write a sequence of strings to the file. The sequence can be any
243        iterable object producing strings, typically a list of strings. There
244        is no return value.
245
246        (The name is intended to match readlines(); writelines() does not add
247        line separators.)
248        """
249        write = self.write
250        for line in iterable:
251            write(line)
252
253    def flush(self):
254        """Flush the internal buffer
255        """
256        _complain_ifclosed(self.closed)
257
258    def getvalue(self):
259        """
260        Retrieve the entire contents of the "file" at any time before
261        the StringIO object's close() method is called.
262
263        The StringIO object can accept either Unicode or 8-bit strings,
264        but mixing the two may take some care. If both are used, 8-bit
265        strings that cannot be interpreted as 7-bit ASCII (that use the
266        8th bit) will cause a UnicodeError to be raised when getvalue()
267        is called.
268        """
269        _complain_ifclosed(self.closed)
270        if self.buflist:
271            self.buf += ''.join(self.buflist)
272            self.buflist = []
273        return self.buf
274
275
276# A little test suite
277
278def test():
279    import sys
280    if sys.argv[1:]:
281        file = sys.argv[1]
282    else:
283        file = '/etc/passwd'
284    lines = open(file, 'r').readlines()
285    text = open(file, 'r').read()
286    f = StringIO()
287    for line in lines[:-2]:
288        f.write(line)
289    f.writelines(lines[-2:])
290    if f.getvalue() != text:
291        raise RuntimeError, 'write failed'
292    length = f.tell()
293    print 'File length =', length
294    f.seek(len(lines[0]))
295    f.write(lines[1])
296    f.seek(0)
297    print 'First line =', repr(f.readline())
298    print 'Position =', f.tell()
299    line = f.readline()
300    print 'Second line =', repr(line)
301    f.seek(-len(line), 1)
302    line2 = f.read(len(line))
303    if line != line2:
304        raise RuntimeError, 'bad result after seek back'
305    f.seek(len(line2), 1)
306    list = f.readlines()
307    line = list[-1]
308    f.seek(f.tell() - len(line))
309    line2 = f.read()
310    if line != line2:
311        raise RuntimeError, 'bad result after seek back from EOF'
312    print 'Read', len(list), 'more lines'
313    print 'File length =', f.tell()
314    if f.tell() != length:
315        raise RuntimeError, 'bad length'
316    f.truncate(length/2)
317    f.seek(0, 2)
318    print 'Truncated length =', f.tell()
319    if f.tell() != length/2:
320        raise RuntimeError, 'truncate did not adjust length'
321    f.close()
322
323if __name__ == '__main__':
324    test()
325