StringIO

Module StringIO |
[hide private]
[frames] | no frames]

# Source Code for Module StringIO

      1  r"""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 -def _complain_ifclosed(closed):
     39      if closed:
     40          raise ValueError, "I/O operation on closed file"
     41
     42 -class StringIO:
     43      """class StringIO([buffer])
     44
     45      When a StringIO object is created, it can be initialized to an existing
     46      string by passing the string to the constructor. If no string is given,
     47      the StringIO will start empty.
     48
     49      The StringIO object can accept either Unicode or 8-bit strings, but
     50      mixing the two may take some care. If both are used, 8-bit strings that
     51      cannot be interpreted as 7-bit ASCII (that use the 8th bit) will cause
     52      a UnicodeError to be raised when getvalue() is called.
     53      """
     54 -    def __init__(self, buf = ''):
     55          # Force self.buf to be a string or unicode
     56          if not isinstance(buf, basestring):
     57              buf = str(buf)
     58          self.buf = buf
     59          self.len = len(buf)
     60          self.buflist = []
     61          self.pos = 0
     62          self.closed = False
     63          self.softspace = 0
     64
     65 -    def __iter__(self):
     66          return self
     67
     68 -    def next(self):
     69          """A file object is its own iterator, for example iter(f) returns f
     70          (unless f is closed). When a file is used as an iterator, typically
     71          in a for loop (for example, for line in f: print line), the next()
     72          method is called repeatedly. This method returns the next input line,
     73          or raises StopIteration when EOF is hit.
     74          """
     75          if self.closed:
     76              raise StopIteration
     77          r = self.readline()
     78          if not r:
     79              raise StopIteration
     80          return r
     81
     82 -    def close(self):
     83          """Free the memory buffer.
     84          """
     85          if not self.closed:
     86              self.closed = True
     87              del self.buf, self.pos
     88
     89 -    def isatty(self):
     90          """Returns False because StringIO objects are not connected to a
     91          tty-like device.
     92          """
     93          _complain_ifclosed(self.closed)
     94          return False
     95
     96 -    def seek(self, pos, mode = 0):
     97          """Set the file's current position.
     98
     99          The mode argument is optional and defaults to 0 (absolute file
    100          positioning); other values are 1 (seek relative to the current
    101          position) and 2 (seek relative to the file's end).
    102
    103          There is no return value.
    104          """
    105          _complain_ifclosed(self.closed)
    106          if self.buflist:
    107              self.buf += ''.join(self.buflist)
    108              self.buflist = []
    109          if mode == 1:
    110              pos += self.pos
    111          elif mode == 2:
    112              pos += self.len
    113          self.pos = max(0, pos)
    114
    115 -    def tell(self):
    116          """Return the file's current position."""
    117          _complain_ifclosed(self.closed)
    118          return self.pos
    119
    120 -    def read(self, n = -1):
    121          """Read at most size bytes from the file
    122          (less if the read hits EOF before obtaining size bytes).
    123
    124          If the size argument is negative or omitted, read all data until EOF
    125          is reached. The bytes are returned as a string object. An empty
    126          string is returned when EOF is encountered immediately.
    127          """
    128          _complain_ifclosed(self.closed)
    129          if self.buflist:
    130              self.buf += ''.join(self.buflist)
    131              self.buflist = []
    132          if n < 0:
    133              newpos = self.len
    134          else:
    135              newpos = min(self.pos+n, self.len)
    136          r = self.buf[self.pos:newpos]
    137          self.pos = newpos
    138          return r
    139
    140 -    def readline(self, length=None):
    141          """Read one entire line from the file.
    142
    143          A trailing newline character is kept in the string (but may be absent
    144          when a file ends with an incomplete line). If the size argument is
    145          present and non-negative, it is a maximum byte count (including the
    146          trailing newline) and an incomplete line may be returned.
    147
    148          An empty string is returned only when EOF is encountered immediately.
    149
    150          Note: Unlike stdio's fgets(), the returned string contains null
    151          characters ('\0') if they occurred in the input.
    152          """
    153          _complain_ifclosed(self.closed)
    154          if self.buflist:
    155              self.buf += ''.join(self.buflist)
    156              self.buflist = []
    157          i = self.buf.find('\n', self.pos)
    158          if i < 0:
    159              newpos = self.len
    160          else:
    161              newpos = i+1
    162          if length is not None:
    163              if self.pos + length < newpos:
    164                  newpos = self.pos + length
    165          r = self.buf[self.pos:newpos]
    166          self.pos = newpos
    167          return r
    168
    169 -    def readlines(self, sizehint = 0):
    170          """Read until EOF using readline() and return a list containing the
    171          lines thus read.
    172
    173          If the optional sizehint argument is present, instead of reading up
    174          to EOF, whole lines totalling approximately sizehint bytes (or more
    175          to accommodate a final whole line).
    176          """
    177          total = 0
    178          lines = []
    179          line = self.readline()
    180          while line:
    181              lines.append(line)
    182              total += len(line)
    183              if 0 < sizehint <= total:
    184                  break
    185              line = self.readline()
    186          return lines
    187
    188 -    def truncate(self, size=None):
    189          """Truncate the file's size.
    190
    191          If the optional size argument is present, the file is truncated to
    192          (at most) that size. The size defaults to the current position.
    193          The current file position is not changed unless the position
    194          is beyond the new file size.
    195
    196          If the specified size exceeds the file's current size, the
    197          file remains unchanged.
    198          """
    199          _complain_ifclosed(self.closed)
    200          if size is None:
    201              size = self.pos
    202          elif size < 0:
    203              raise IOError(EINVAL, "Negative size not allowed")
    204          elif size < self.pos:
    205              self.pos = size
    206          self.buf = self.getvalue()[:size]
    207          self.len = size
    208
    209 -    def write(self, s):
    210          """Write a string to the file.
    211
    212          There is no return value.
    213          """
    214          _complain_ifclosed(self.closed)
    215          if not s: return
    216          # Force s to be a string or unicode
    217          if not isinstance(s, basestring):
    218              s = str(s)
    219          spos = self.pos
    220          slen = self.len
    221          if spos == slen:
    222              self.buflist.append(s)
    223              self.len = self.pos = spos + len(s)
    224              return
    225          if spos > slen:
    226              self.buflist.append('\0'*(spos - slen))
    227              slen = spos
    228          newpos = spos + len(s)
    229          if spos < slen:
    230              if self.buflist:
    231                  self.buf += ''.join(self.buflist)
    232              self.buflist = [self.buf[:spos], s, self.buf[newpos:]]
    233              self.buf = ''
    234              if newpos > slen:
    235                  slen = newpos
    236          else:
    237              self.buflist.append(s)
    238              slen = newpos
    239          self.len = slen
    240          self.pos = newpos
    241
    242 -    def writelines(self, iterable):
    243          """Write a sequence of strings to the file. The sequence can be any
    244          iterable object producing strings, typically a list of strings. There
    245          is no return value.
    246
    247          (The name is intended to match readlines(); writelines() does not add
    248          line separators.)
    249          """
    250          write = self.write
    251          for line in iterable:
    252              write(line)
    253
    254 -    def flush(self):
    255          """Flush the internal buffer
    256          """
    257          _complain_ifclosed(self.closed)
    258
    259 -    def getvalue(self):
    260          """
    261          Retrieve the entire contents of the "file" at any time before
    262          the StringIO object's close() method is called.
    263
    264          The StringIO object can accept either Unicode or 8-bit strings,
    265          but mixing the two may take some care. If both are used, 8-bit
    266          strings that cannot be interpreted as 7-bit ASCII (that use the
    267          8th bit) will cause a UnicodeError to be raised when getvalue()
    268          is called.
    269          """
    270          if self.buflist:
    271              self.buf += ''.join(self.buflist)
    272              self.buflist = []
    273          return self.buf
    274
    275
    276  # A little test suite
    277
    278 -def test():
    279      import sys
    280      if sys.argv[1:]:
    281          file = sys.argv[1]
    282      else:
    283          file = '/etc/passwd'
    284      lines = open(file, 'r').readlines()
    285      text = open(file, 'r').read()
    286      f = StringIO()
    287      for line in lines[:-2]:
    288          f.write(line)
    289      f.writelines(lines[-2:])
    290      if f.getvalue() != text:
    291          raise RuntimeError, 'write failed'
    292      length = f.tell()
    293      print 'File length =', length
    294      f.seek(len(lines[0]))
    295      f.write(lines[1])
    296      f.seek(0)
    297      print 'First line =', repr(f.readline())
    298      print 'Position =', f.tell()
    299      line = f.readline()
    300      print 'Second line =', repr(line)
    301      f.seek(-len(line), 1)
    302      line2 = f.read(len(line))
    303      if line != line2:
    304          raise RuntimeError, 'bad result after seek back'
    305      f.seek(len(line2), 1)
    306      list = f.readlines()
    307      line = list[-1]
    308      f.seek(f.tell() - len(line))
    309      line2 = f.read()
    310      if line != line2:
    311          raise RuntimeError, 'bad result after seek back from EOF'
    312      print 'Read', len(list), 'more lines'
    313      print 'File length =', f.tell()
    314      if f.tell() != length:
    315          raise RuntimeError, 'bad length'
    316      f.truncate(length/2)
    317      f.seek(0, 2)
    318      print 'Truncated length =', f.tell()
    319      if f.tell() != length/2:
    320          raise RuntimeError, 'truncate did not adjust length'
    321      f.close()
    322
    323  if __name__ == '__main__':
    324      test()
    325