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