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