Lib/StringIO.py

   1 r"""File-like objects that read from or write to a string buffer.
   2
   3 This implements (nearly) all stdio methods.
   4
   5 f = StringIO()      # ready for writing
   6 f = StringIO(buf)   # ready for reading
   7 f.close()           # explicitly release resources held
   8 flag = f.isatty()   # always false
   9 pos = f.tell()      # get current position
  10 f.seek(pos)         # set current position
  11 f.seek(pos, mode)   # mode 0: absolute; 1: relative; 2: relative to EOF
  12 buf = f.read()      # read until EOF
  13 buf = f.read(n)     # read up to n bytes
  14 buf = f.readline()  # read until end of line ('\n') or EOF
  15 list = f.readlines()# list of f.readline() results until EOF
  16 f.truncate([size])  # truncate file at to at most size (default: current pos)
  17 f.write(buf)        # write at current position
  18 f.writelines(list)  # for line in list: f.write(line)
  19 f.getvalue()        # return whole file's contents as a string
  20
  21 Notes:
  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 """
  31 try:
  32     from errno import EINVAL
  33 except ImportError:
  34     EINVAL = 22
  35
  36 __all__ = ["StringIO"]
  37
  38 def _complain_ifclosed(closed):
  39     if closed:
  40         raise ValueError, "I/O operation on closed file"
  41
  42 class 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         if self.closed:
  76             raise StopIteration
  77         r = self.readline()
  78         if not r:
  79             raise StopIteration
  80         return r
  81
  82     def close(self):
  83         """Free the memory buffer.
  84         """
  85         if not self.closed:
  86             self.closed = True
  87             del self.buf, self.pos
  88
  89     def isatty(self):
  90         """Returns False because StringIO objects are not connected to a
  91         tty-like device.
  92         """
  93         _complain_ifclosed(self.closed)
  94         return False
  95
  96     def seek(self, pos, mode = 0):
  97         """Set the file's current position.
  98
  99         The mode argument is optional and defaults to 0 (absolute file
 100         positioning); other values are 1 (seek relative to the current
 101         position) and 2 (seek relative to the file's end).
 102
 103         There is no return value.
 104         """
 105         _complain_ifclosed(self.closed)
 106         if self.buflist:
 107             self.buf += ''.join(self.buflist)
 108             self.buflist = []
 109         if mode == 1:
 110             pos += self.pos
 111         elif mode == 2:
 112             pos += self.len
 113         self.pos = max(0, pos)
 114
 115     def tell(self):
 116         """Return the file's current position."""
 117         _complain_ifclosed(self.closed)
 118         return self.pos
 119
 120     def read(self, n = -1):
 121         """Read at most size bytes from the file
 122         (less if the read hits EOF before obtaining size bytes).
 123
 124         If the size argument is negative or omitted, read all data until EOF
 125         is reached. The bytes are returned as a string object. An empty
 126         string is returned when EOF is encountered immediately.
 127         """
 128         _complain_ifclosed(self.closed)
 129         if self.buflist:
 130             self.buf += ''.join(self.buflist)
 131             self.buflist = []
 132         if n < 0:
 133             newpos = self.len
 134         else:
 135             newpos = min(self.pos+n, self.len)
 136         r = self.buf[self.pos:newpos]
 137         self.pos = newpos
 138         return r
 139
 140     def readline(self, length=None):
 141         """Read one entire line from the file.
 142
 143         A trailing newline character is kept in the string (but may be absent
 144         when a file ends with an incomplete line). If the size argument is
 145         present and non-negative, it is a maximum byte count (including the
 146         trailing newline) and an incomplete line may be returned.
 147
 148         An empty string is returned only when EOF is encountered immediately.
 149
 150         Note: Unlike stdio's fgets(), the returned string contains null
 151         characters ('\0') if they occurred in the input.
 152         """
 153         _complain_ifclosed(self.closed)
 154         if self.buflist:
 155             self.buf += ''.join(self.buflist)
 156             self.buflist = []
 157         i = self.buf.find('\n', self.pos)
 158         if i < 0:
 159             newpos = self.len
 160         else:
 161             newpos = i+1
 162         if length is not None:
 163             if self.pos + length < newpos:
 164                 newpos = self.pos + length
 165         r = self.buf[self.pos:newpos]
 166         self.pos = newpos
 167         return r
 168
 169     def readlines(self, sizehint = 0):
 170         """Read until EOF using readline() and return a list containing the
 171         lines thus read.
 172
 173         If the optional sizehint argument is present, instead of reading up
 174         to EOF, whole lines totalling approximately sizehint bytes (or more
 175         to accommodate a final whole line).
 176         """
 177         total = 0
 178         lines = []
 179         line = self.readline()
 180         while line:
 181             lines.append(line)
 182             total += len(line)
 183             if 0 < sizehint <= total:
 184                 break
 185             line = self.readline()
 186         return lines
 187
 188     def truncate(self, size=None):
 189         """Truncate the file's size.
 190
 191         If the optional size argument is present, the file is truncated to
 192         (at most) that size. The size defaults to the current position.
 193         The current file position is not changed unless the position
 194         is beyond the new file size.
 195
 196         If the specified size exceeds the file's current size, the
 197         file remains unchanged.
 198         """
 199         _complain_ifclosed(self.closed)
 200         if size is None:
 201             size = self.pos
 202         elif size < 0:
 203             raise IOError(EINVAL, "Negative size not allowed")
 204         elif size < self.pos:
 205             self.pos = size
 206         self.buf = self.getvalue()[:size]
 207         self.len = size
 208
 209     def write(self, s):
 210         """Write a string to the file.
 211
 212         There is no return value.
 213         """
 214         _complain_ifclosed(self.closed)
 215         if not s: return
 216         # Force s to be a string or unicode
 217         if not isinstance(s, basestring):
 218             s = str(s)
 219         spos = self.pos
 220         slen = self.len
 221         if spos == slen:
 222             self.buflist.append(s)
 223             self.len = self.pos = spos + len(s)
 224             return
 225         if spos > slen:
 226             self.buflist.append('\0'*(spos - slen))
 227             slen = spos
 228         newpos = spos + len(s)
 229         if spos < slen:
 230             if self.buflist:
 231                 self.buf += ''.join(self.buflist)
 232             self.buflist = [self.buf[:spos], s, self.buf[newpos:]]
 233             self.buf = ''
 234             if newpos > slen:
 235                 slen = newpos
 236         else:
 237             self.buflist.append(s)
 238             slen = newpos
 239         self.len = slen
 240         self.pos = newpos
 241
 242     def writelines(self, iterable):
 243         """Write a sequence of strings to the file. The sequence can be any
 244         iterable object producing strings, typically a list of strings. There
 245         is no return value.
 246
 247         (The name is intended to match readlines(); writelines() does not add
 248         line separators.)
 249         """
 250         write = self.write
 251         for line in iterable:
 252             write(line)
 253
 254     def flush(self):
 255         """Flush the internal buffer
 256         """
 257         _complain_ifclosed(self.closed)
 258
 259     def getvalue(self):
 260         """
 261         Retrieve the entire contents of the "file" at any time before
 262         the StringIO object's close() method is called.
 263
 264         The StringIO object can accept either Unicode or 8-bit strings,
 265         but mixing the two may take some care. If both are used, 8-bit
 266         strings that cannot be interpreted as 7-bit ASCII (that use the
 267         8th bit) will cause a UnicodeError to be raised when getvalue()
 268         is called.
 269         """
 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
 278 def 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
 323 if __name__ == '__main__':
 324     test()