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.pos = 0
  35                 self.closed = 0
  36                 self.softspace = 0
  37         def close(self):
  38                 if not self.closed:
  39                         self.closed = 1
  40                         del self.buf, self.pos
  41         def isatty(self):
  42                 return 0
  43         def seek(self, pos, mode = 0):
  44                 if mode == 1:
  45                         pos = pos + self.pos
  46                 elif mode == 2:
  47                         pos = pos + len(self.buf)
  48                 self.pos = max(0, pos)
  49         def tell(self):
  50                 return self.pos
  51         def read(self, n = 0):
  52                 if n <= 0:
  53                         newpos = len(self.buf)
  54                 else:
  55                         newpos = min(self.pos+n, len(self.buf))
  56                 r = self.buf[self.pos:newpos]
  57                 self.pos = newpos
  58                 return r
  59         def readline(self):
  60                 i = string.find(self.buf, '\n', self.pos)
  61                 if i < 0:
  62                         newpos = len(self.buf)
  63                 else:
  64                         newpos = i+1
  65                 r = self.buf[self.pos:newpos]
  66                 self.pos = newpos
  67                 return r
  68         def readlines(self):
  69                 lines = []
  70                 line = self.readline()
  71                 while line:
  72                         lines.append(line)
  73                         line = self.readline()
  74                 return lines
  75         def write(self, s):
  76                 if not s: return
  77                 if self.pos > len(self.buf):
  78                         self.buf = self.buf + '\0'*(self.pos - len(self.buf))
  79                 newpos = self.pos + len(s)
  80                 self.buf = self.buf[:self.pos] + s + self.buf[newpos:]
  81                 self.pos = newpos
  82         def writelines(self, list):
  83                 self.write(string.joinfields(list, ''))
  84         def getvalue(self):
  85                 return self.buf
  86
  87
  88 # A little test suite
  89
  90 def test():
  91         import sys
  92         if sys.argv[1:]:
  93                 file = sys.argv[1]
  94         else:
  95                 file = '/etc/passwd'
  96         lines = open(file, 'r').readlines()
  97         text = open(file, 'r').read()
  98         f = StringIO()
  99         for line in lines[:-2]:
 100                 f.write(line)
 101         f.writelines(lines[-2:])
 102         if f.getvalue() != text:
 103                 raise RuntimeError, 'write failed'
 104         length = f.tell()
 105         print 'File length =', length
 106         f.seek(len(lines[0]))
 107         f.write(lines[1])
 108         f.seek(0)
 109         print 'First line =', `f.readline()`
 110         here = f.tell()
 111         line = f.readline()
 112         print 'Second line =', `line`
 113         f.seek(-len(line), 1)
 114         line2 = f.read(len(line))
 115         if line != line2:
 116                 raise RuntimeError, 'bad result after seek back'
 117         f.seek(len(line2), 1)
 118         list = f.readlines()
 119         line = list[-1]
 120         f.seek(f.tell() - len(line))
 121         line2 = f.read()
 122         if line != line2:
 123                 raise RuntimeError, 'bad result after seek back from EOF'
 124         print 'Read', len(list), 'more lines'
 125         print 'File length =', f.tell()
 126         if f.tell() != length:
 127                 raise RuntimeError, 'bad length'
 128         f.close()
 129
 130 if __name__ == '__main__':
 131         test()