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