blob: 8efd7d8c48f9f733a3c6ad74c4b9905b2019ce0f [file] [log] [blame]
Guido van Rossumaad67612000-05-08 17:31:04 +00001"""File-like objects that read from or write to a string buffer.
2
3This implements (nearly) all stdio methods.
4
5f = StringIO() # ready for writing
6f = StringIO(buf) # ready for reading
7f.close() # explicitly release resources held
8flag = f.isatty() # always false
9pos = f.tell() # get current position
10f.seek(pos) # set current position
11f.seek(pos, mode) # mode 0: absolute; 1: relative; 2: relative to EOF
12buf = f.read() # read until EOF
13buf = f.read(n) # read up to n bytes
14buf = f.readline() # read until end of line ('\n') or EOF
15list = f.readlines()# list of f.readline() results until EOF
16f.write(buf) # write at current position
17f.writelines(list) # for line in list: f.write(line)
18f.getvalue() # return whole file's contents as a string
19
20Notes:
21- Using a real file is often faster (but less convenient).
22- There's also a much faster implementation in C, called cStringIO, but
23 it's not subclassable.
24- fileno() is left unimplemented so that code which uses it triggers
25 an exception early.
26- Seeking far beyond EOF and then writing will insert real null
27 bytes that occupy space in the buffer.
28- There's a simple test set (see end of this file).
29"""
Guido van Rossum5c971671996-07-22 15:23:25 +000030
31import string
32
33class StringIO:
34 def __init__(self, buf = ''):
35 self.buf = buf
36 self.len = len(buf)
37 self.buflist = []
38 self.pos = 0
39 self.closed = 0
40 self.softspace = 0
41 def close(self):
42 if not self.closed:
43 self.closed = 1
44 del self.buf, self.pos
45 def isatty(self):
Guido van Rossum64e736b1998-10-02 01:23:47 +000046 if self.closed:
47 raise ValueError, "I/O operation on closed file"
Guido van Rossum5c971671996-07-22 15:23:25 +000048 return 0
49 def seek(self, pos, mode = 0):
Guido van Rossum64e736b1998-10-02 01:23:47 +000050 if self.closed:
51 raise ValueError, "I/O operation on closed file"
Guido van Rossum5c971671996-07-22 15:23:25 +000052 if self.buflist:
53 self.buf = self.buf + string.joinfields(self.buflist, '')
54 self.buflist = []
55 if mode == 1:
56 pos = pos + self.pos
57 elif mode == 2:
58 pos = pos + self.len
59 self.pos = max(0, pos)
60 def tell(self):
Guido van Rossum64e736b1998-10-02 01:23:47 +000061 if self.closed:
62 raise ValueError, "I/O operation on closed file"
Guido van Rossum5c971671996-07-22 15:23:25 +000063 return self.pos
64 def read(self, n = -1):
Guido van Rossum64e736b1998-10-02 01:23:47 +000065 if self.closed:
66 raise ValueError, "I/O operation on closed file"
Guido van Rossum5c971671996-07-22 15:23:25 +000067 if self.buflist:
68 self.buf = self.buf + string.joinfields(self.buflist, '')
69 self.buflist = []
70 if n < 0:
71 newpos = self.len
72 else:
73 newpos = min(self.pos+n, self.len)
74 r = self.buf[self.pos:newpos]
75 self.pos = newpos
76 return r
Guido van Rossumff712aa1997-08-15 00:45:26 +000077 def readline(self, length=None):
Guido van Rossum64e736b1998-10-02 01:23:47 +000078 if self.closed:
79 raise ValueError, "I/O operation on closed file"
Guido van Rossum5c971671996-07-22 15:23:25 +000080 if self.buflist:
81 self.buf = self.buf + string.joinfields(self.buflist, '')
82 self.buflist = []
83 i = string.find(self.buf, '\n', self.pos)
84 if i < 0:
85 newpos = self.len
86 else:
87 newpos = i+1
Guido van Rossumff712aa1997-08-15 00:45:26 +000088 if length is not None:
89 if self.pos + length < newpos:
90 newpos = self.pos + length
Guido van Rossum5c971671996-07-22 15:23:25 +000091 r = self.buf[self.pos:newpos]
92 self.pos = newpos
93 return r
94 def readlines(self):
95 lines = []
96 line = self.readline()
97 while line:
98 lines.append(line)
99 line = self.readline()
100 return lines
101 def write(self, s):
Guido van Rossum64e736b1998-10-02 01:23:47 +0000102 if self.closed:
103 raise ValueError, "I/O operation on closed file"
Guido van Rossum5c971671996-07-22 15:23:25 +0000104 if not s: return
105 if self.pos > self.len:
106 self.buflist.append('\0'*(self.pos - self.len))
107 self.len = self.pos
108 newpos = self.pos + len(s)
109 if self.pos < self.len:
110 if self.buflist:
111 self.buf = self.buf + string.joinfields(self.buflist, '')
112 self.buflist = []
113 self.buflist = [self.buf[:self.pos], s, self.buf[newpos:]]
114 self.buf = ''
115 else:
116 self.buflist.append(s)
117 self.len = newpos
118 self.pos = newpos
119 def writelines(self, list):
120 self.write(string.joinfields(list, ''))
121 def flush(self):
Guido van Rossum64e736b1998-10-02 01:23:47 +0000122 if self.closed:
123 raise ValueError, "I/O operation on closed file"
Guido van Rossum5c971671996-07-22 15:23:25 +0000124 def getvalue(self):
125 if self.buflist:
126 self.buf = self.buf + string.joinfields(self.buflist, '')
127 self.buflist = []
128 return self.buf
129
130
131# A little test suite
132
133def test():
134 import sys
135 if sys.argv[1:]:
136 file = sys.argv[1]
137 else:
138 file = '/etc/passwd'
139 lines = open(file, 'r').readlines()
140 text = open(file, 'r').read()
141 f = StringIO()
142 for line in lines[:-2]:
143 f.write(line)
144 f.writelines(lines[-2:])
145 if f.getvalue() != text:
146 raise RuntimeError, 'write failed'
147 length = f.tell()
148 print 'File length =', length
149 f.seek(len(lines[0]))
150 f.write(lines[1])
151 f.seek(0)
152 print 'First line =', `f.readline()`
153 here = f.tell()
154 line = f.readline()
155 print 'Second line =', `line`
156 f.seek(-len(line), 1)
157 line2 = f.read(len(line))
158 if line != line2:
159 raise RuntimeError, 'bad result after seek back'
160 f.seek(len(line2), 1)
161 list = f.readlines()
162 line = list[-1]
163 f.seek(f.tell() - len(line))
164 line2 = f.read()
165 if line != line2:
166 raise RuntimeError, 'bad result after seek back from EOF'
167 print 'Read', len(list), 'more lines'
168 print 'File length =', f.tell()
169 if f.tell() != length:
170 raise RuntimeError, 'bad length'
171 f.close()
172
173if __name__ == '__main__':
174 test()