Completed
Push — master ( 1a32f0...6baa7b )
by Klaus
01:15
created

TeeingStreamProxy.__init__()   A

Complexity

Conditions 1

Size

Total Lines 3

Duplication

Lines 0
Ratio 0 %

Importance

Changes 1
Bugs 0 Features 0
Metric Value
cc 1
dl 0
loc 3
rs 10
c 1
b 0
f 0
1
#!/usr/bin/env python
2
# coding=utf-8
3
from __future__ import division, print_function, unicode_literals
4
import os
5
import sys
6
from io import StringIO
7
import subprocess
8
from threading import Timer
9
from contextlib import contextmanager
10
import wrapt
11
from sacred.optional import libc
12
from tempfile import NamedTemporaryFile
13
from sacred.settings import SETTINGS
14
15
__sacred__ = True  # marks files that should be filtered from stack traces
16
17
18
# A PY2 compatible FileNotFoundError
19
if sys.version_info[0] == 2:
20
    import errno
21
22
    class FileNotFoundError(IOError):
23
        def __init__(self, msg):
24
            super(FileNotFoundError, self).__init__(errno.ENOENT, msg)
25
else:
26
    # Reassign so that we can import it from here
27
    FileNotFoundError = FileNotFoundError
28
29
30
def flush():
31
    """Try to flush all stdio buffers, both from python and from C."""
32
    try:
33
        sys.stdout.flush()
34
        sys.stderr.flush()
35
    except (AttributeError, ValueError, IOError):
36
        pass  # unsupported
37
    try:
38
        libc.fflush(None)
39
    except (AttributeError, ValueError, IOError):
40
        pass  # unsupported
41
42
43
def get_stdcapturer(mode=None):
44
    mode = mode if mode is not None else SETTINGS.STDOUT_CAPTURING
45
    return {
46
        None: no_tee,
47
        "None": no_tee,
48
        "FD": tee_output_fd,
49
        "PY": tee_output_python}[mode]
50
51
52
class TeeingStreamProxy(wrapt.ObjectProxy):
53
    """A wrapper around stdout or stderr that duplicates all output to out."""
54
55
    def __init__(self, wrapped, out):
56
        super(TeeingStreamProxy, self).__init__(wrapped)
57
        self._self_out = out
58
59
    def write(self, data):
60
        self.__wrapped__.write(data)
61
        self._self_out.write(data)
62
63
    def flush(self):
64
        self.__wrapped__.flush()
65
        self._self_out.flush()
66
67
68
@contextmanager
69
def no_tee():
70
    out = StringIO()
71
    final_out = [""]
72
    yield out, final_out
73
    out.close()
74
75
76
@contextmanager
77
def tee_output_python():
78
    """Duplicate sys.stdout and sys.stderr to new StringIO."""
79
    out = StringIO()
80
    final_out = []
81
    flush()
82
    orig_stdout, orig_stderr = sys.stdout, sys.stderr
83
    sys.stdout = TeeingStreamProxy(sys.stdout, out)
84
    sys.stderr = TeeingStreamProxy(sys.stderr, out)
85
    yield out, final_out
86
    flush()
87
    sys.stdout, sys.stderr = orig_stdout, orig_stderr
88
    final_out.append(out.getvalue())
89
    out.close()
90
91
92
# Duplicate stdout and stderr to a file. Inspired by:
93
# http://eli.thegreenplace.net/2015/redirecting-all-kinds-of-stdout-in-python/
94
# http://stackoverflow.com/a/651718/1388435
95
# http://stackoverflow.com/a/22434262/1388435
96
@contextmanager
97
def tee_output_fd():
98
    """Duplicate stdout and stderr to a file on the file descriptor level."""
99
    with NamedTemporaryFile() as target:
100
        original_stdout_fd = 1
101
        original_stderr_fd = 2
102
103
        # Save a copy of the original stdout and stderr file descriptors
104
        saved_stdout_fd = os.dup(original_stdout_fd)
105
        saved_stderr_fd = os.dup(original_stderr_fd)
106
107
        target_fd = target.fileno()
108
109
        final_output = []
110
111
        try:
112
            tee_stdout = subprocess.Popen(
113
                ['tee', '-a', '/dev/stderr'],
114
                stdin=subprocess.PIPE, stderr=target_fd, stdout=1)
115
            tee_stderr = subprocess.Popen(
116
                ['tee', '-a', '/dev/stderr'],
117
                stdin=subprocess.PIPE, stderr=target_fd, stdout=2)
118
        except (FileNotFoundError, OSError):
119
            tee_stdout = subprocess.Popen(
120
                [sys.executable, "-m", "sacred.pytee"],
121
                stdin=subprocess.PIPE, stderr=target_fd)
122
            tee_stderr = subprocess.Popen(
123
                [sys.executable, "-m", "sacred.pytee"],
124
                stdin=subprocess.PIPE, stdout=target_fd)
125
126
        flush()
127
        os.dup2(tee_stdout.stdin.fileno(), original_stdout_fd)
128
        os.dup2(tee_stderr.stdin.fileno(), original_stderr_fd)
129
130
        try:
131
            yield target, final_output  # let the caller do their printing
132
        finally:
133
            flush()
134
135
            # then redirect stdout back to the saved fd
136
            tee_stdout.stdin.close()
137
            tee_stderr.stdin.close()
138
139
            # restore original fds
140
            os.dup2(saved_stdout_fd, original_stdout_fd)
141
            os.dup2(saved_stderr_fd, original_stderr_fd)
142
143
            # wait for completion of the tee processes with timeout
144
            # implemented using a timer because timeout support is py3 only
145
            def kill_tees():
146
                tee_stdout.kill()
147
                tee_stderr.kill()
148
149
            tee_timer = Timer(1, kill_tees)
150
            try:
151
                tee_timer.start()
152
                tee_stdout.wait()
153
                tee_stderr.wait()
154
            finally:
155
                tee_timer.cancel()
156
157
            os.close(saved_stdout_fd)
158
            os.close(saved_stderr_fd)
159
            target.flush()
160
            target.seek(0)
161
            final_output.append(target.read().decode())
162