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