Coverage for sacred/sacred/stdout_capturing.py: 25%
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
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
1#!/usr/bin/env python
2# coding=utf-8
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
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
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]
41class TeeingStreamProxy(wrapt.ObjectProxy):
42 """A wrapper around stdout or stderr that duplicates all output to out."""
44 def __init__(self, wrapped, out):
45 super().__init__(wrapped)
46 self._self_out = out
48 def write(self, data):
49 self.__wrapped__.write(data)
50 self._self_out.write(data)
52 def flush(self):
53 self.__wrapped__.flush()
54 self._self_out.flush()
57class CapturedStdout:
58 def __init__(self, buffer):
59 self.buffer = buffer
60 self.read_position = 0
61 self.final = None
63 @property
64 def closed(self):
65 return self.buffer.closed
67 def flush(self):
68 return self.buffer.flush()
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
81 def finalize(self):
82 self.flush()
83 self.final = self.get()
84 self.buffer.close()
87@contextmanager
88def no_tee():
89 out = CapturedStdout(StringIO())
90 try:
91 yield out
92 finally:
93 out.finalize()
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
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()
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)
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 )
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)
163 try:
164 yield out # let the caller do their printing
165 finally:
166 flush()
168 # then redirect stdout back to the saved fd
169 tee_stdout.stdin.close()
170 tee_stderr.stdin.close()
172 # restore original fds
173 os.dup2(saved_stdout_fd, original_stdout_fd)
174 os.dup2(saved_stderr_fd, original_stderr_fd)
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()
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()
188 os.close(saved_stdout_fd)
189 os.close(saved_stderr_fd)
190 out.finalize()