def captured(stdout=CaptureTarget.STRING, stderr=CaptureTarget.STRING):
"""Capture outputs of sys.stdout and sys.stderr.
If stdout is STRING, capture sys.stdout as a string,
if stdout is None, do not capture sys.stdout, leaving it untouched,
otherwise redirect sys.stdout to the file-like object given by stdout.
Behave correspondingly for stderr with the exception that if stderr is STDOUT,
redirect sys.stderr to stdout target and set stderr attribute of yielded object to None.
Args:
stdout: capture target for sys.stdout, one of STRING, None, or file-like object
stderr: capture target for sys.stderr, one of STRING, STDOUT, None, or file-like object
Yields:
CapturedText: has attributes stdout, stderr which are either strings, None or the
corresponding file-like function argument.
"""
# NOTE: This function is not thread-safe. Using within multi-threading may cause spurious
# behavior of not returning sys.stdout and sys.stderr back to their 'proper' state
# """
# Context manager to capture the printed output of the code in the with block
#
# Bind the context manager to a variable using `as` and the result will be
# in the stdout property.
#
# >>> from conda.common.io import captured
# >>> with captured() as c:
# ... print('hello world!')
# ...
# >>> c.stdout
# 'hello world!\n'
# """
def write_wrapper(self, to_write):
# This may have to deal with a *lot* of text.
if hasattr(self, 'mode') and 'b' in self.mode:
wanted = bytes
elif isinstance(self, BytesIO):
wanted = bytes
else:
wanted = str
if not isinstance(to_write, wanted):
if hasattr(to_write, 'decode'):
decoded = to_write.decode('utf-8')
self.old_write(decoded)
elif hasattr(to_write, 'encode'):
b = to_write.encode('utf-8')
self.old_write(b)
else:
self.old_write(to_write)
class CapturedText(object):
pass
# sys.stdout.write(u'unicode out')
# sys.stdout.write(bytes('bytes out', encoding='utf-8'))
# sys.stdout.write(str('str out'))
saved_stdout, saved_stderr = sys.stdout, sys.stderr
if stdout == CaptureTarget.STRING:
outfile = StringIO()
outfile.old_write = outfile.write
outfile.write = partial(write_wrapper, outfile)
sys.stdout = outfile
else:
outfile = stdout
if outfile is not None:
sys.stdout = outfile
if stderr == CaptureTarget.STRING:
errfile = StringIO()
errfile.old_write = errfile.write
errfile.write = partial(write_wrapper, errfile)
sys.stderr = errfile
elif stderr == CaptureTarget.STDOUT:
sys.stderr = errfile = outfile
else:
errfile = stderr
if errfile is not None:
sys.stderr = errfile
c = CapturedText()
log.debug("overtaking stderr and stdout")
try:
yield c
finally:
if stdout == CaptureTarget.STRING:
c.stdout = outfile.getvalue()
else:
c.stdout = outfile
if stderr == CaptureTarget.STRING:
c.stderr = errfile.getvalue()
elif stderr == CaptureTarget.STDOUT:
c.stderr = None
else:
c.stderr = errfile
sys.stdout, sys.stderr = saved_stdout, saved_stderr
log.debug("stderr and stdout yielding back")
