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