Lib/StringIO.py

   1 # class StringIO implements  file-like objects that read/write a
   2 # string buffer (a.k.a. "memory files").
   3 #
   4 # This implements (nearly) all stdio methods.
   5 #
   6 # f = StringIO()      # ready for writing
   7 # f = StringIO(buf)   # ready for reading
   8 # f.close()           # explicitly release resources held
   9 # flag = f.isatty()   # always false
  10 # pos = f.tell()      # get current position
  11 # f.seek(pos)         # set current position
  12 # f.seek(pos, mode)   # mode 0: absolute; 1: relative; 2: relative to EOF
  13 # buf = f.read()      # read until EOF
  14 # buf = f.read(n)     # read up to n bytes
  15 # buf = f.readline()  # read until end of line ('\n') or EOF
  16 # list = f.readlines()# list of f.readline() results until EOF
  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 # - fileno() is left unimplemented so that code which uses it triggers
  24 #   an exception early.
  25 # - Seeking far beyond EOF and then writing will insert real null
  26 #   bytes that occupy space in the buffer.
  27 # - There's a simple test set (see end of this file).
  28
  29 import string
  30
  31 class StringIO:
  32         def __init__(self, buf = ''):
  33                 self.buf = buf
  34                 self.len = len(buf)
  35                 self.buflist = []
  36                 self.pos = 0
  37                 self.closed = 0
  38                 self.softspace = 0
  39         def close(self):
  40                 if not self.closed:
  41                         self.closed = 1
  42                         del self.buf, self.pos
  43         def isatty(self):
  44                 if self.closed:
  45                         raise ValueError, "I/O operation on closed file"
  46                 return 0
  47         def seek(self, pos, mode = 0):
  48                 if self.closed:
  49                         raise ValueError, "I/O operation on closed file"
  50                 if self.buflist:
  51                         self.buf = self.buf + string.joinfields(self.buflist, '')
  52                         self.buflist = []
  53                 if mode == 1:
  54                         pos = pos + self.pos
  55                 elif mode == 2:
  56                         pos = pos + self.len
  57                 self.pos = max(0, pos)
  58         def tell(self):
  59                 if self.closed:
  60                         raise ValueError, "I/O operation on closed file"
  61                 return self.pos
  62         def read(self, n = -1):
  63                 if self.closed:
  64                         raise ValueError, "I/O operation on closed file"
  65                 if self.buflist:
  66                         self.buf = self.buf + string.joinfields(self.buflist, '')
  67                         self.buflist = []
  68                 if n < 0:
  69                         newpos = self.len
  70                 else:
  71                         newpos = min(self.pos+n, self.len)
  72                 r = self.buf[self.pos:newpos]
  73                 self.pos = newpos
  74                 return r
  75         def readline(self, length=None):
  76                 if self.closed:
  77                         raise ValueError, "I/O operation on closed file"
  78                 if self.buflist:
  79                         self.buf = self.buf + string.joinfields(self.buflist, '')
  80                         self.buflist = []
  81                 i = string.find(self.buf, '\n', self.pos)
  82                 if i < 0:
  83                         newpos = self.len
  84                 else:
  85                         newpos = i+1
  86                 if length is not None:
  87                         if self.pos + length < newpos:
  88                                 newpos = self.pos + length
  89                 r = self.buf[self.pos:newpos]
  90                 self.pos = newpos
  91                 return r
  92         def readlines(self):
  93                 lines = []
  94                 line = self.readline()
  95                 while line:
  96                         lines.append(line)
  97                         line = self.readline()
  98                 return lines
  99         def write(self, s):
 100                 if self.closed:
 101                         raise ValueError, "I/O operation on closed file"
 102                 if not s: return
 103                 if self.pos > self.len:
 104                         self.buflist.append('\0'*(self.pos - self.len))
 105                         self.len = self.pos
 106                 newpos = self.pos + len(s)
 107                 if self.pos < self.len:
 108                         if self.buflist:
 109                                 self.buf = self.buf + string.joinfields(self.buflist, '')
 110                                 self.buflist = []
 111                         self.buflist = [self.buf[:self.pos], s, self.buf[newpos:]]
 112                         self.buf = ''
 113                 else:
 114                         self.buflist.append(s)
 115                         self.len = newpos
 116                 self.pos = newpos
 117         def writelines(self, list):
 118                 self.write(string.joinfields(list, ''))
 119         def flush(self):
 120                 if self.closed:
 121                         raise ValueError, "I/O operation on closed file"
 122         def getvalue(self):
 123                 if self.buflist:
 124                         self.buf = self.buf + string.joinfields(self.buflist, '')
 125                         self.buflist = []
 126                 return self.buf
 127
 128
 129 # A little test suite
 130
 131 def test():
 132         import sys
 133         if sys.argv[1:]:
 134                 file = sys.argv[1]
 135         else:
 136                 file = '/etc/passwd'
 137         lines = open(file, 'r').readlines()
 138         text = open(file, 'r').read()
 139         f = StringIO()
 140         for line in lines[:-2]:
 141                 f.write(line)
 142         f.writelines(lines[-2:])
 143         if f.getvalue() != text:
 144                 raise RuntimeError, 'write failed'
 145         length = f.tell()
 146         print 'File length =', length
 147         f.seek(len(lines[0]))
 148         f.write(lines[1])
 149         f.seek(0)
 150         print 'First line =', `f.readline()`
 151         here = f.tell()
 152         line = f.readline()
 153         print 'Second line =', `line`
 154         f.seek(-len(line), 1)
 155         line2 = f.read(len(line))
 156         if line != line2:
 157                 raise RuntimeError, 'bad result after seek back'
 158         f.seek(len(line2), 1)
 159         list = f.readlines()
 160         line = list[-1]
 161         f.seek(f.tell() - len(line))
 162         line2 = f.read()
 163         if line != line2:
 164                 raise RuntimeError, 'bad result after seek back from EOF'
 165         print 'Read', len(list), 'more lines'
 166         print 'File length =', f.tell()
 167         if f.tell() != length:
 168                 raise RuntimeError, 'bad length'
 169         f.close()
 170
 171 if __name__ == '__main__':
 172         test()