Lib/StringIO.py

   1 """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 class StringIO:
  39     """class StringIO([buffer])
  40
  41     When a StringIO object is created, it can be initialized to an existing
  42     string by passing the string to the constructor. If no string is given,
  43     the StringIO will start empty.
  44
  45     The StringIO object can accept either Unicode or 8-bit strings, but
  46     mixing the two may take some care. If both are used, 8-bit strings that
  47     cannot be interpreted as 7-bit ASCII (that use the 8th bit) will cause
  48     a UnicodeError to be raised when getvalue() is called.
  49     """
  50     def __init__(self, buf = ''):
  51         # Force self.buf to be a string or unicode
  52         if not isinstance(buf, basestring):
  53             buf = str(buf)
  54         self.buf = buf
  55         self.len = len(buf)
  56         self.buflist = []
  57         self.pos = 0
  58         self.closed = 0
  59         self.softspace = 0
  60
  61     def __iter__(self):
  62         return self
  63
  64     def next(self):
  65         if self.closed:
  66             raise StopIteration
  67         r = self.readline()
  68         if not r:
  69             raise StopIteration
  70         return r
  71
  72     def close(self):
  73         """Free the memory buffer.
  74         """
  75         if not self.closed:
  76             self.closed = 1
  77             del self.buf, self.pos
  78
  79     def isatty(self):
  80         if self.closed:
  81             raise ValueError, "I/O operation on closed file"
  82         return False
  83
  84     def seek(self, pos, mode = 0):
  85         if self.closed:
  86             raise ValueError, "I/O operation on closed file"
  87         if self.buflist:
  88             self.buf += ''.join(self.buflist)
  89             self.buflist = []
  90         if mode == 1:
  91             pos += self.pos
  92         elif mode == 2:
  93             pos += self.len
  94         self.pos = max(0, pos)
  95
  96     def tell(self):
  97         if self.closed:
  98             raise ValueError, "I/O operation on closed file"
  99         return self.pos
 100
 101     def read(self, n = -1):
 102         if self.closed:
 103             raise ValueError, "I/O operation on closed file"
 104         if self.buflist:
 105             self.buf += ''.join(self.buflist)
 106             self.buflist = []
 107         if n < 0:
 108             newpos = self.len
 109         else:
 110             newpos = min(self.pos+n, self.len)
 111         r = self.buf[self.pos:newpos]
 112         self.pos = newpos
 113         return r
 114
 115     def readline(self, length=None):
 116         if self.closed:
 117             raise ValueError, "I/O operation on closed file"
 118         if self.buflist:
 119             self.buf += ''.join(self.buflist)
 120             self.buflist = []
 121         i = self.buf.find('\n', self.pos)
 122         if i < 0:
 123             newpos = self.len
 124         else:
 125             newpos = i+1
 126         if length is not None:
 127             if self.pos + length < newpos:
 128                 newpos = self.pos + length
 129         r = self.buf[self.pos:newpos]
 130         self.pos = newpos
 131         return r
 132
 133     def readlines(self, sizehint = 0):
 134         total = 0
 135         lines = []
 136         line = self.readline()
 137         while line:
 138             lines.append(line)
 139             total += len(line)
 140             if 0 < sizehint <= total:
 141                 break
 142             line = self.readline()
 143         return lines
 144
 145     def truncate(self, size=None):
 146         if self.closed:
 147             raise ValueError, "I/O operation on closed file"
 148         if size is None:
 149             size = self.pos
 150         elif size < 0:
 151             raise IOError(EINVAL, "Negative size not allowed")
 152         elif size < self.pos:
 153             self.pos = size
 154         self.buf = self.getvalue()[:size]
 155
 156     def write(self, s):
 157         if self.closed:
 158             raise ValueError, "I/O operation on closed file"
 159         if not s: return
 160         # Force s to be a string or unicode
 161         if not isinstance(s, basestring):
 162             s = str(s)
 163         if self.pos == self.len:
 164             self.buflist.append(s)
 165             self.len = self.pos = self.pos + len(s)
 166             return
 167         if self.pos > self.len:
 168             self.buflist.append('\0'*(self.pos - self.len))
 169             self.len = self.pos
 170         newpos = self.pos + len(s)
 171         if self.pos < self.len:
 172             if self.buflist:
 173                 self.buf += ''.join(self.buflist)
 174                 self.buflist = []
 175             self.buflist = [self.buf[:self.pos], s, self.buf[newpos:]]
 176             self.buf = ''
 177             if newpos > self.len:
 178                 self.len = newpos
 179         else:
 180             self.buflist.append(s)
 181             self.len = newpos
 182         self.pos = newpos
 183
 184     def writelines(self, list):
 185         self.write(''.join(list))
 186
 187     def flush(self):
 188         if self.closed:
 189             raise ValueError, "I/O operation on closed file"
 190
 191     def getvalue(self):
 192         """
 193         Retrieve the entire contents of the "file" at any time before
 194         the StringIO object's close() method is called.
 195
 196         The StringIO object can accept either Unicode or 8-bit strings,
 197         but mixing the two may take some care. If both are used, 8-bit
 198         strings that cannot be interpreted as 7-bit ASCII (that use the
 199         8th bit) will cause a UnicodeError to be raised when getvalue()
 200         is called.
 201         """
 202         if self.buflist:
 203             self.buf += ''.join(self.buflist)
 204             self.buflist = []
 205         return self.buf
 206
 207
 208 # A little test suite
 209
 210 def test():
 211     import sys
 212     if sys.argv[1:]:
 213         file = sys.argv[1]
 214     else:
 215         file = '/etc/passwd'
 216     lines = open(file, 'r').readlines()
 217     text = open(file, 'r').read()
 218     f = StringIO()
 219     for line in lines[:-2]:
 220         f.write(line)
 221     f.writelines(lines[-2:])
 222     if f.getvalue() != text:
 223         raise RuntimeError, 'write failed'
 224     length = f.tell()
 225     print 'File length =', length
 226     f.seek(len(lines[0]))
 227     f.write(lines[1])
 228     f.seek(0)
 229     print 'First line =', `f.readline()`
 230     print 'Position =', f.tell()
 231     line = f.readline()
 232     print 'Second line =', `line`
 233     f.seek(-len(line), 1)
 234     line2 = f.read(len(line))
 235     if line != line2:
 236         raise RuntimeError, 'bad result after seek back'
 237     f.seek(len(line2), 1)
 238     list = f.readlines()
 239     line = list[-1]
 240     f.seek(f.tell() - len(line))
 241     line2 = f.read()
 242     if line != line2:
 243         raise RuntimeError, 'bad result after seek back from EOF'
 244     print 'Read', len(list), 'more lines'
 245     print 'File length =', f.tell()
 246     if f.tell() != length:
 247         raise RuntimeError, 'bad length'
 248     f.close()
 249
 250 if __name__ == '__main__':
 251     test()