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
  32 try:
  33     from errno import EINVAL
  34 except ImportError:
  35     EINVAL = 22
  36
  37 __all__ = ["StringIO"]
  38
  39 class StringIO:
  40     def __init__(self, buf = ''):
  41         # Force self.buf to be a string
  42         self.buf = str(buf)
  43         self.len = len(buf)
  44         self.buflist = []
  45         self.pos = 0
  46         self.closed = 0
  47         self.softspace = 0
  48
  49     def __iter__(self):
  50         return iter(self.readline, '')
  51
  52     def close(self):
  53         if not self.closed:
  54             self.closed = 1
  55             del self.buf, self.pos
  56
  57     def isatty(self):
  58         if self.closed:
  59             raise ValueError, "I/O operation on closed file"
  60         return 0
  61
  62     def seek(self, pos, mode = 0):
  63         if self.closed:
  64             raise ValueError, "I/O operation on closed file"
  65         if self.buflist:
  66             self.buf += ''.join(self.buflist)
  67             self.buflist = []
  68         if mode == 1:
  69             pos += self.pos
  70         elif mode == 2:
  71             pos += self.len
  72         self.pos = max(0, pos)
  73
  74     def tell(self):
  75         if self.closed:
  76             raise ValueError, "I/O operation on closed file"
  77         return self.pos
  78
  79     def read(self, n = -1):
  80         if self.closed:
  81             raise ValueError, "I/O operation on closed file"
  82         if self.buflist:
  83             self.buf += ''.join(self.buflist)
  84             self.buflist = []
  85         if n < 0:
  86             newpos = self.len
  87         else:
  88             newpos = min(self.pos+n, self.len)
  89         r = self.buf[self.pos:newpos]
  90         self.pos = newpos
  91         return r
  92
  93     def readline(self, length=None):
  94         if self.closed:
  95             raise ValueError, "I/O operation on closed file"
  96         if self.buflist:
  97             self.buf += ''.join(self.buflist)
  98             self.buflist = []
  99         i = self.buf.find('\n', self.pos)
 100         if i < 0:
 101             newpos = self.len
 102         else:
 103             newpos = i+1
 104         if length is not None:
 105             if self.pos + length < newpos:
 106                 newpos = self.pos + length
 107         r = self.buf[self.pos:newpos]
 108         self.pos = newpos
 109         return r
 110
 111     def readlines(self, sizehint = 0):
 112         total = 0
 113         lines = []
 114         line = self.readline()
 115         while line:
 116             lines.append(line)
 117             total += len(line)
 118             if 0 < sizehint <= total:
 119                 break
 120             line = self.readline()
 121         return lines
 122
 123     def truncate(self, size=None):
 124         if self.closed:
 125             raise ValueError, "I/O operation on closed file"
 126         if size is None:
 127             size = self.pos
 128         elif size < 0:
 129             raise IOError(EINVAL, "Negative size not allowed")
 130         elif size < self.pos:
 131             self.pos = size
 132         self.buf = self.getvalue()[:size]
 133
 134     def write(self, s):
 135         if self.closed:
 136             raise ValueError, "I/O operation on closed file"
 137         if not s: return
 138         # Force s to be a string
 139         s = str(s)
 140         if self.pos > self.len:
 141             self.buflist.append('\0'*(self.pos - self.len))
 142             self.len = self.pos
 143         newpos = self.pos + len(s)
 144         if self.pos < self.len:
 145             if self.buflist:
 146                 self.buf += ''.join(self.buflist)
 147                 self.buflist = []
 148             self.buflist = [self.buf[:self.pos], s, self.buf[newpos:]]
 149             self.buf = ''
 150             if newpos > self.len:
 151                 self.len = newpos
 152         else:
 153             self.buflist.append(s)
 154             self.len = newpos
 155         self.pos = newpos
 156
 157     def writelines(self, list):
 158         self.write(''.join(list))
 159
 160     def flush(self):
 161         if self.closed:
 162             raise ValueError, "I/O operation on closed file"
 163
 164     def getvalue(self):
 165         if self.buflist:
 166             self.buf += ''.join(self.buflist)
 167             self.buflist = []
 168         return self.buf
 169
 170
 171 # A little test suite
 172
 173 def test():
 174     import sys
 175     if sys.argv[1:]:
 176         file = sys.argv[1]
 177     else:
 178         file = '/etc/passwd'
 179     lines = open(file, 'r').readlines()
 180     text = open(file, 'r').read()
 181     f = StringIO()
 182     for line in lines[:-2]:
 183         f.write(line)
 184     f.writelines(lines[-2:])
 185     if f.getvalue() != text:
 186         raise RuntimeError, 'write failed'
 187     length = f.tell()
 188     print 'File length =', length
 189     f.seek(len(lines[0]))
 190     f.write(lines[1])
 191     f.seek(0)
 192     print 'First line =', `f.readline()`
 193     here = f.tell()
 194     line = f.readline()
 195     print 'Second line =', `line`
 196     f.seek(-len(line), 1)
 197     line2 = f.read(len(line))
 198     if line != line2:
 199         raise RuntimeError, 'bad result after seek back'
 200     f.seek(len(line2), 1)
 201     list = f.readlines()
 202     line = list[-1]
 203     f.seek(f.tell() - len(line))
 204     line2 = f.read()
 205     if line != line2:
 206         raise RuntimeError, 'bad result after seek back from EOF'
 207     print 'Read', len(list), 'more lines'
 208     print 'File length =', f.tell()
 209     if f.tell() != length:
 210         raise RuntimeError, 'bad length'
 211     f.close()
 212
 213 if __name__ == '__main__':
 214     test()