1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
|
"""
Test decorator for capturing stdout/stderr/both.
Based on original code from Fabric 1.x, specifically:
* fabric/tests/utils.py
* as of Git SHA 62abc4e17aab0124bf41f9c5f9c4bc86cc7d9412
Though modifications have been made since.
"""
import io
import sys
from functools import wraps
class CarbonCopy(io.BytesIO):
"""
An IO wrapper capable of multiplexing its writes to other buffer objects.
"""
def __init__(self, buffer=b"", cc=None):
"""
If ``cc`` is given and is a file-like object or an iterable of same,
it/they will be written to whenever this instance is written to.
"""
super().__init__(buffer)
if cc is None:
cc = []
elif hasattr(cc, "write"):
cc = [cc]
self.cc = cc
def write(self, s):
# Ensure we always write bytes.
if isinstance(s, str):
s = s.encode("utf-8")
# Write out to our capturing object & any CC's
super().write(s)
for writer in self.cc:
writer.write(s)
# Real sys.std(out|err) requires writing to a buffer attribute obj in some
# situations.
@property
def buffer(self):
return self
# Make sure we always hand back strings
def getvalue(self):
ret = super().getvalue()
if isinstance(ret, bytes):
ret = ret.decode("utf-8")
return ret
def trap(func):
"""
Replace sys.std(out|err) with a wrapper during execution, restored after.
In addition, a new combined-streams output (another wrapper) will appear at
``sys.stdall``. This stream will resemble what a user sees at a terminal,
i.e. both out/err streams intermingled.
"""
@wraps(func)
def wrapper(*args, **kwargs):
# Use another CarbonCopy even though we're not cc'ing; for our "write
# bytes, return strings" behavior. Meh.
sys.stdall = CarbonCopy()
my_stdout, sys.stdout = sys.stdout, CarbonCopy(cc=sys.stdall)
my_stderr, sys.stderr = sys.stderr, CarbonCopy(cc=sys.stdall)
try:
return func(*args, **kwargs)
finally:
sys.stdout = my_stdout
sys.stderr = my_stderr
del sys.stdall
return wrapper
|
