Coverage for sacred/sacred/stdout_capturing.py: 88%

Shortcuts on this page

r m x   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

116 statements  

1#!/usr/bin/env python 

2# coding=utf-8 

3 

4import os 

5import sys 

6import subprocess 

7import warnings 

8from io import StringIO 

9from contextlib import contextmanager 

10import wrapt 

11from sacred.optional import libc 

12from tempfile import NamedTemporaryFile 

13from sacred.settings import SETTINGS 

14 

15 

16def flush(): 

17 """Try to flush all stdio buffers, both from python and from C.""" 

18 try: 

19 sys.stdout.flush() 

20 sys.stderr.flush() 

21 except (AttributeError, ValueError, OSError): 

22 pass # unsupported 

23 try: 

24 libc.fflush(None) 

25 except (AttributeError, ValueError, OSError): 

26 pass # unsupported 

27 

28 

29def get_stdcapturer(mode=None): 

30 mode = mode if mode is not None else SETTINGS.CAPTURE_MODE 

31 capture_options = {"no": no_tee, "fd": tee_output_fd, "sys": tee_output_python} 

32 if mode not in capture_options: 

33 raise KeyError( 

34 "Unknown capture mode '{}'. Available options are {}".format( 

35 mode, sorted(capture_options.keys()) 

36 ) 

37 ) 

38 return mode, capture_options[mode] 

39 

40 

41class TeeingStreamProxy(wrapt.ObjectProxy): 

42 """A wrapper around stdout or stderr that duplicates all output to out.""" 

43 

44 def __init__(self, wrapped, out): 

45 super().__init__(wrapped) 

46 self._self_out = out 

47 

48 def write(self, data): 

49 self.__wrapped__.write(data) 

50 self._self_out.write(data) 

51 

52 def flush(self): 

53 self.__wrapped__.flush() 

54 self._self_out.flush() 

55 

56 

57class CapturedStdout: 

58 def __init__(self, buffer): 

59 self.buffer = buffer 

60 self.read_position = 0 

61 self.final = None 

62 

63 @property 

64 def closed(self): 

65 return self.buffer.closed 

66 

67 def flush(self): 

68 return self.buffer.flush() 

69 

70 def get(self): 

71 if self.final is None: 

72 self.buffer.seek(self.read_position) 

73 value = self.buffer.read() 

74 self.read_position = self.buffer.tell() 

75 return value 

76 else: 

77 value = self.final 

78 self.final = None 

79 return value 

80 

81 def finalize(self): 

82 self.flush() 

83 self.final = self.get() 

84 self.buffer.close() 

85 

86 

87@contextmanager 

88def no_tee(): 

89 out = CapturedStdout(StringIO()) 

90 try: 

91 yield out 

92 finally: 

93 out.finalize() 

94 

95 

96@contextmanager 

97def tee_output_python(): 

98 """Duplicate sys.stdout and sys.stderr to new StringIO.""" 

99 buffer = StringIO() 

100 out = CapturedStdout(buffer) 

101 orig_stdout, orig_stderr = sys.stdout, sys.stderr 

102 flush() 

103 sys.stdout = TeeingStreamProxy(sys.stdout, buffer) 

104 sys.stderr = TeeingStreamProxy(sys.stderr, buffer) 

105 try: 

106 yield out 

107 finally: 

108 flush() 

109 out.finalize() 

110 sys.stdout, sys.stderr = orig_stdout, orig_stderr 

111 

112 

113# Duplicate stdout and stderr to a file. Inspired by: 

114# http://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/ 

115# http://stackoverflow.com/a/651718/1388435 

116# http://stackoverflow.com/a/22434262/1388435 

117@contextmanager 

118def tee_output_fd(): 

119 """Duplicate stdout and stderr to a file on the file descriptor level.""" 

120 with NamedTemporaryFile(mode="w+", newline="") as target: 

121 original_stdout_fd = 1 

122 original_stderr_fd = 2 

123 target_fd = target.fileno() 

124 

125 # Save a copy of the original stdout and stderr file descriptors) 

126 saved_stdout_fd = os.dup(original_stdout_fd) 

127 saved_stderr_fd = os.dup(original_stderr_fd) 

128 

129 try: 

130 # start_new_session=True to move process to a new process group 

131 # this is done to avoid receiving KeyboardInterrupts (see #149) 

132 tee_stdout = subprocess.Popen( 

133 ["tee", "-a", target.name], 

134 start_new_session=True, 

135 stdin=subprocess.PIPE, 

136 stdout=1, 

137 ) 

138 tee_stderr = subprocess.Popen( 

139 ["tee", "-a", target.name], 

140 start_new_session=True, 

141 stdin=subprocess.PIPE, 

142 stdout=2, 

143 ) 

144 except (FileNotFoundError, OSError, AttributeError): 

145 # No tee found in this operating system. Trying to use a python 

146 # implementation of tee. However this is slow and error-prone. 

147 tee_stdout = subprocess.Popen( 

148 [sys.executable, "-m", "sacred.pytee"], 

149 stdin=subprocess.PIPE, 

150 stderr=target_fd, 

151 ) 

152 tee_stderr = subprocess.Popen( 

153 [sys.executable, "-m", "sacred.pytee"], 

154 stdin=subprocess.PIPE, 

155 stdout=target_fd, 

156 ) 

157 

158 flush() 

159 os.dup2(tee_stdout.stdin.fileno(), original_stdout_fd) 

160 os.dup2(tee_stderr.stdin.fileno(), original_stderr_fd) 

161 out = CapturedStdout(target) 

162 

163 try: 

164 yield out # let the caller do their printing 

165 finally: 

166 flush() 

167 

168 # then redirect stdout back to the saved fd 

169 tee_stdout.stdin.close() 

170 tee_stderr.stdin.close() 

171 

172 # restore original fds 

173 os.dup2(saved_stdout_fd, original_stdout_fd) 

174 os.dup2(saved_stderr_fd, original_stderr_fd) 

175 

176 try: 

177 tee_stdout.wait(timeout=1) 

178 except subprocess.TimeoutExpired: 

179 warnings.warn("tee_stdout.wait timeout. Forcibly terminating.") 

180 tee_stdout.terminate() 

181 

182 try: 

183 tee_stderr.wait(timeout=1) 

184 except subprocess.TimeoutExpired: 

185 warnings.warn("tee_stderr.wait timeout. Forcibly terminating.") 

186 tee_stderr.terminate() 

187 

188 os.close(saved_stdout_fd) 

189 os.close(saved_stderr_fd) 

190 out.finalize()