PageRenderTime 362ms CodeModel.GetById 141ms app.highlight 84ms RepoModel.GetById 117ms app.codeStats 0ms

/Lib/StringIO.py

http://unladen-swallow.googlecode.com/
Python | 323 lines | 323 code | 0 blank | 0 comment | 1 complexity | 7f5407f1548ef3aa47f2044fdfcbca6a MD5 | raw file
  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 < 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:
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        if self.buflist:
270            self.buf += ''.join(self.buflist)
271            self.buflist = []
272        return self.buf
273
274
275# A little test suite
276
277def test():
278    import sys
279    if sys.argv[1:]:
280        file = sys.argv[1]
281    else:
282        file = '/etc/passwd'
283    lines = open(file, 'r').readlines()
284    text = open(file, 'r').read()
285    f = StringIO()
286    for line in lines[:-2]:
287        f.write(line)
288    f.writelines(lines[-2:])
289    if f.getvalue() != text:
290        raise RuntimeError, 'write failed'
291    length = f.tell()
292    print 'File length =', length
293    f.seek(len(lines[0]))
294    f.write(lines[1])
295    f.seek(0)
296    print 'First line =', repr(f.readline())
297    print 'Position =', f.tell()
298    line = f.readline()
299    print 'Second line =', repr(line)
300    f.seek(-len(line), 1)
301    line2 = f.read(len(line))
302    if line != line2:
303        raise RuntimeError, 'bad result after seek back'
304    f.seek(len(line2), 1)
305    list = f.readlines()
306    line = list[-1]
307    f.seek(f.tell() - len(line))
308    line2 = f.read()
309    if line != line2:
310        raise RuntimeError, 'bad result after seek back from EOF'
311    print 'Read', len(list), 'more lines'
312    print 'File length =', f.tell()
313    if f.tell() != length:
314        raise RuntimeError, 'bad length'
315    f.truncate(length/2)
316    f.seek(0, 2)
317    print 'Truncated length =', f.tell()
318    if f.tell() != length/2:
319        raise RuntimeError, 'truncate did not adjust length'
320    f.close()
321
322if __name__ == '__main__':
323    test()