blob: ca8c49308daa6ef5544901c3eb68665e9b0ccb2b [file] [log] [blame]
Tor Norbye3a2425a2013-11-04 10:16:08 -08001class IORedirector:
2 '''This class works to redirect the write function to many streams
3 '''
4
5 def __init__(self, *args):
6 self._redirectTo = args
7
8 def write(self, s):
9 for r in self._redirectTo:
10 try:
11 r.write(s)
12 except:
13 pass
14
15 def isatty(self):
16 return False #not really a file
17
18 def flush(self):
19 for r in self._redirectTo:
20 r.flush()
21
22 def __getattr__(self, name):
23 for r in self._redirectTo:
24 if hasattr(r, name):
25 return r.__getattribute__(name)
26 raise AttributeError(name)
27
28class IOBuf:
29 '''This class works as a replacement for stdio and stderr.
30 It is a buffer and when its contents are requested, it will erase what
31
32 it has so far so that the next return will not return the same contents again.
33 '''
34 def __init__(self):
35 self.buflist = []
36
37 def getvalue(self):
38 b = self.buflist
39 self.buflist = [] #clear it
40 return ''.join(b)
41
42 def write(self, s):
43 self.buflist.append(s)
44
45 def isatty(self):
46 return False #not really a file
47
48 def flush(self):
49 pass
50
51
52class _RedirectionsHolder:
53 _stack_stdout = []
54 _stack_stderr = []
55
56
57def StartRedirect(keep_original_redirection=False, std='stdout'):
58 '''
59 @param std: 'stdout', 'stderr', or 'both'
60 '''
61 import sys
62 buf = IOBuf()
63
64 if std == 'both':
65 config_stds = ['stdout', 'stderr']
66 else:
67 config_stds = [std]
68
69 for std in config_stds:
70 original = getattr(sys, std)
71 stack = getattr(_RedirectionsHolder, '_stack_%s' % std)
72 stack.append(original)
73
74 if keep_original_redirection:
75 setattr(sys, std, IORedirector(buf, getattr(sys, std)))
76 else:
77 setattr(sys, std, buf)
78 return buf
79
80
81def EndRedirect(std='stdout'):
82 import sys
83 if std == 'both':
84 config_stds = ['stdout', 'stderr']
85 else:
86 config_stds = [std]
87 for std in config_stds:
88 stack = getattr(_RedirectionsHolder, '_stack_%s' % std)
89 setattr(sys, std, stack.pop())
90
91