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")