##// END OF EJS Templates
str(CapturedIO) gives stdout
MinRK -
Show More
@@ -1,381 +1,384
1 1 # encoding: utf-8
2 2 """
3 3 IO related utilities.
4 4 """
5 5
6 6 #-----------------------------------------------------------------------------
7 7 # Copyright (C) 2008-2011 The IPython Development Team
8 8 #
9 9 # Distributed under the terms of the BSD License. The full license is in
10 10 # the file COPYING, distributed as part of this software.
11 11 #-----------------------------------------------------------------------------
12 12 from __future__ import print_function
13 13
14 14 #-----------------------------------------------------------------------------
15 15 # Imports
16 16 #-----------------------------------------------------------------------------
17 17 import os
18 18 import sys
19 19 import tempfile
20 20 from StringIO import StringIO
21 21
22 22 #-----------------------------------------------------------------------------
23 23 # Code
24 24 #-----------------------------------------------------------------------------
25 25
26 26
27 27 class IOStream:
28 28
29 29 def __init__(self,stream, fallback=None):
30 30 if not hasattr(stream,'write') or not hasattr(stream,'flush'):
31 31 if fallback is not None:
32 32 stream = fallback
33 33 else:
34 34 raise ValueError("fallback required, but not specified")
35 35 self.stream = stream
36 36 self._swrite = stream.write
37 37
38 38 # clone all methods not overridden:
39 39 def clone(meth):
40 40 return not hasattr(self, meth) and not meth.startswith('_')
41 41 for meth in filter(clone, dir(stream)):
42 42 setattr(self, meth, getattr(stream, meth))
43 43
44 44 def write(self,data):
45 45 try:
46 46 self._swrite(data)
47 47 except:
48 48 try:
49 49 # print handles some unicode issues which may trip a plain
50 50 # write() call. Emulate write() by using an empty end
51 51 # argument.
52 52 print(data, end='', file=self.stream)
53 53 except:
54 54 # if we get here, something is seriously broken.
55 55 print('ERROR - failed to write data to stream:', self.stream,
56 56 file=sys.stderr)
57 57
58 58 def writelines(self, lines):
59 59 if isinstance(lines, basestring):
60 60 lines = [lines]
61 61 for line in lines:
62 62 self.write(line)
63 63
64 64 # This class used to have a writeln method, but regular files and streams
65 65 # in Python don't have this method. We need to keep this completely
66 66 # compatible so we removed it.
67 67
68 68 @property
69 69 def closed(self):
70 70 return self.stream.closed
71 71
72 72 def close(self):
73 73 pass
74 74
75 75 # setup stdin/stdout/stderr to sys.stdin/sys.stdout/sys.stderr
76 76 devnull = open(os.devnull, 'a')
77 77 stdin = IOStream(sys.stdin, fallback=devnull)
78 78 stdout = IOStream(sys.stdout, fallback=devnull)
79 79 stderr = IOStream(sys.stderr, fallback=devnull)
80 80
81 81 class IOTerm:
82 82 """ Term holds the file or file-like objects for handling I/O operations.
83 83
84 84 These are normally just sys.stdin, sys.stdout and sys.stderr but for
85 85 Windows they can can replaced to allow editing the strings before they are
86 86 displayed."""
87 87
88 88 # In the future, having IPython channel all its I/O operations through
89 89 # this class will make it easier to embed it into other environments which
90 90 # are not a normal terminal (such as a GUI-based shell)
91 91 def __init__(self, stdin=None, stdout=None, stderr=None):
92 92 mymodule = sys.modules[__name__]
93 93 self.stdin = IOStream(stdin, mymodule.stdin)
94 94 self.stdout = IOStream(stdout, mymodule.stdout)
95 95 self.stderr = IOStream(stderr, mymodule.stderr)
96 96
97 97
98 98 class Tee(object):
99 99 """A class to duplicate an output stream to stdout/err.
100 100
101 101 This works in a manner very similar to the Unix 'tee' command.
102 102
103 103 When the object is closed or deleted, it closes the original file given to
104 104 it for duplication.
105 105 """
106 106 # Inspired by:
107 107 # http://mail.python.org/pipermail/python-list/2007-May/442737.html
108 108
109 109 def __init__(self, file_or_name, mode="w", channel='stdout'):
110 110 """Construct a new Tee object.
111 111
112 112 Parameters
113 113 ----------
114 114 file_or_name : filename or open filehandle (writable)
115 115 File that will be duplicated
116 116
117 117 mode : optional, valid mode for open().
118 118 If a filename was give, open with this mode.
119 119
120 120 channel : str, one of ['stdout', 'stderr']
121 121 """
122 122 if channel not in ['stdout', 'stderr']:
123 123 raise ValueError('Invalid channel spec %s' % channel)
124 124
125 125 if hasattr(file_or_name, 'write') and hasattr(file_or_name, 'seek'):
126 126 self.file = file_or_name
127 127 else:
128 128 self.file = open(file_or_name, mode)
129 129 self.channel = channel
130 130 self.ostream = getattr(sys, channel)
131 131 setattr(sys, channel, self)
132 132 self._closed = False
133 133
134 134 def close(self):
135 135 """Close the file and restore the channel."""
136 136 self.flush()
137 137 setattr(sys, self.channel, self.ostream)
138 138 self.file.close()
139 139 self._closed = True
140 140
141 141 def write(self, data):
142 142 """Write data to both channels."""
143 143 self.file.write(data)
144 144 self.ostream.write(data)
145 145 self.ostream.flush()
146 146
147 147 def flush(self):
148 148 """Flush both channels."""
149 149 self.file.flush()
150 150 self.ostream.flush()
151 151
152 152 def __del__(self):
153 153 if not self._closed:
154 154 self.close()
155 155
156 156
157 157 def file_read(filename):
158 158 """Read a file and close it. Returns the file source."""
159 159 fobj = open(filename,'r');
160 160 source = fobj.read();
161 161 fobj.close()
162 162 return source
163 163
164 164
165 165 def file_readlines(filename):
166 166 """Read a file and close it. Returns the file source using readlines()."""
167 167 fobj = open(filename,'r');
168 168 lines = fobj.readlines();
169 169 fobj.close()
170 170 return lines
171 171
172 172
173 173 def raw_input_multi(header='', ps1='==> ', ps2='..> ',terminate_str = '.'):
174 174 """Take multiple lines of input.
175 175
176 176 A list with each line of input as a separate element is returned when a
177 177 termination string is entered (defaults to a single '.'). Input can also
178 178 terminate via EOF (^D in Unix, ^Z-RET in Windows).
179 179
180 180 Lines of input which end in \\ are joined into single entries (and a
181 181 secondary continuation prompt is issued as long as the user terminates
182 182 lines with \\). This allows entering very long strings which are still
183 183 meant to be treated as single entities.
184 184 """
185 185
186 186 try:
187 187 if header:
188 188 header += '\n'
189 189 lines = [raw_input(header + ps1)]
190 190 except EOFError:
191 191 return []
192 192 terminate = [terminate_str]
193 193 try:
194 194 while lines[-1:] != terminate:
195 195 new_line = raw_input(ps1)
196 196 while new_line.endswith('\\'):
197 197 new_line = new_line[:-1] + raw_input(ps2)
198 198 lines.append(new_line)
199 199
200 200 return lines[:-1] # don't return the termination command
201 201 except EOFError:
202 202 print()
203 203 return lines
204 204
205 205
206 206 def raw_input_ext(prompt='', ps2='... '):
207 207 """Similar to raw_input(), but accepts extended lines if input ends with \\."""
208 208
209 209 line = raw_input(prompt)
210 210 while line.endswith('\\'):
211 211 line = line[:-1] + raw_input(ps2)
212 212 return line
213 213
214 214
215 215 def ask_yes_no(prompt,default=None):
216 216 """Asks a question and returns a boolean (y/n) answer.
217 217
218 218 If default is given (one of 'y','n'), it is used if the user input is
219 219 empty. Otherwise the question is repeated until an answer is given.
220 220
221 221 An EOF is treated as the default answer. If there is no default, an
222 222 exception is raised to prevent infinite loops.
223 223
224 224 Valid answers are: y/yes/n/no (match is not case sensitive)."""
225 225
226 226 answers = {'y':True,'n':False,'yes':True,'no':False}
227 227 ans = None
228 228 while ans not in answers.keys():
229 229 try:
230 230 ans = raw_input(prompt+' ').lower()
231 231 if not ans: # response was an empty string
232 232 ans = default
233 233 except KeyboardInterrupt:
234 234 pass
235 235 except EOFError:
236 236 if default in answers.keys():
237 237 ans = default
238 238 print()
239 239 else:
240 240 raise
241 241
242 242 return answers[ans]
243 243
244 244
245 245 class NLprinter:
246 246 """Print an arbitrarily nested list, indicating index numbers.
247 247
248 248 An instance of this class called nlprint is available and callable as a
249 249 function.
250 250
251 251 nlprint(list,indent=' ',sep=': ') -> prints indenting each level by 'indent'
252 252 and using 'sep' to separate the index from the value. """
253 253
254 254 def __init__(self):
255 255 self.depth = 0
256 256
257 257 def __call__(self,lst,pos='',**kw):
258 258 """Prints the nested list numbering levels."""
259 259 kw.setdefault('indent',' ')
260 260 kw.setdefault('sep',': ')
261 261 kw.setdefault('start',0)
262 262 kw.setdefault('stop',len(lst))
263 263 # we need to remove start and stop from kw so they don't propagate
264 264 # into a recursive call for a nested list.
265 265 start = kw['start']; del kw['start']
266 266 stop = kw['stop']; del kw['stop']
267 267 if self.depth == 0 and 'header' in kw.keys():
268 268 print(kw['header'])
269 269
270 270 for idx in range(start,stop):
271 271 elem = lst[idx]
272 272 newpos = pos + str(idx)
273 273 if type(elem)==type([]):
274 274 self.depth += 1
275 275 self.__call__(elem, newpos+",", **kw)
276 276 self.depth -= 1
277 277 else:
278 278 print(kw['indent']*self.depth + newpos + kw["sep"] + repr(elem))
279 279
280 280 nlprint = NLprinter()
281 281
282 282
283 283 def temp_pyfile(src, ext='.py'):
284 284 """Make a temporary python file, return filename and filehandle.
285 285
286 286 Parameters
287 287 ----------
288 288 src : string or list of strings (no need for ending newlines if list)
289 289 Source code to be written to the file.
290 290
291 291 ext : optional, string
292 292 Extension for the generated file.
293 293
294 294 Returns
295 295 -------
296 296 (filename, open filehandle)
297 297 It is the caller's responsibility to close the open file and unlink it.
298 298 """
299 299 fname = tempfile.mkstemp(ext)[1]
300 300 f = open(fname,'w')
301 301 f.write(src)
302 302 f.flush()
303 303 return fname, f
304 304
305 305
306 306 def raw_print(*args, **kw):
307 307 """Raw print to sys.__stdout__, otherwise identical interface to print()."""
308 308
309 309 print(*args, sep=kw.get('sep', ' '), end=kw.get('end', '\n'),
310 310 file=sys.__stdout__)
311 311 sys.__stdout__.flush()
312 312
313 313
314 314 def raw_print_err(*args, **kw):
315 315 """Raw print to sys.__stderr__, otherwise identical interface to print()."""
316 316
317 317 print(*args, sep=kw.get('sep', ' '), end=kw.get('end', '\n'),
318 318 file=sys.__stderr__)
319 319 sys.__stderr__.flush()
320 320
321 321
322 322 # Short aliases for quick debugging, do NOT use these in production code.
323 323 rprint = raw_print
324 324 rprinte = raw_print_err
325 325
326 326
327 327 class CapturedIO(object):
328 328 """Simple object for containing captured stdout/err StringIO objects"""
329 329
330 330 def __init__(self, stdout, stderr):
331 331 self._stdout = stdout
332 332 self._stderr = stderr
333 333
334 def __str__(self):
335 return self.stdout
336
334 337 @property
335 338 def stdout(self):
336 339 if not self._stdout:
337 340 return ''
338 341 return self._stdout.getvalue()
339 342
340 343 @property
341 344 def stderr(self):
342 345 if not self._stderr:
343 346 return ''
344 347 return self._stderr.getvalue()
345 348
346 349 def show(self):
347 350 """write my output to sys.stdout/err as appropriate"""
348 351 sys.stdout.write(self.stdout)
349 352 sys.stderr.write(self.stderr)
350 353 sys.stdout.flush()
351 354 sys.stderr.flush()
352 355
353 356 __call__ = show
354 357
355 358
356 359 class capture_output(object):
357 360 """context manager for capturing stdout/err"""
358 361 stdout = True
359 362 stderr = True
360 363
361 364 def __init__(self, stdout=True, stderr=True):
362 365 self.stdout = stdout
363 366 self.stderr = stderr
364 367
365 368 def __enter__(self):
366 369 self.sys_stdout = sys.stdout
367 370 self.sys_stderr = sys.stderr
368 371
369 372 stdout = stderr = False
370 373 if self.stdout:
371 374 stdout = sys.stdout = StringIO()
372 375 if self.stderr:
373 376 stderr = sys.stderr = StringIO()
374 377
375 378 return CapturedIO(stdout, stderr)
376 379
377 380 def __exit__(self, exc_type, exc_value, traceback):
378 381 sys.stdout = self.sys_stdout
379 382 sys.stderr = self.sys_stderr
380 383
381 384
General Comments 0
You need to be logged in to leave comments. Login now