还借鉴了@kindall 和@ForeveWintr 的答案,这是一个完成此任务的类。与先前答案的主要区别在于,这将其捕获为字符串,而不是StringIO
对象,使用起来更方便!
import io
from collections import UserString
from contextlib import redirect_stdout
class capture(UserString, str, redirect_stdout):
'''
Captures stdout (e.g., from ``print()``) as a variable.
Based on ``contextlib.redirect_stdout``, but saves the user the trouble of
defining and reading from an IO stream. Useful for testing the output of functions
that are supposed to print certain output.
'''
def __init__(self, seq='', *args, **kwargs):
self._io = io.StringIO()
UserString.__init__(self, seq=seq, *args, **kwargs)
redirect_stdout.__init__(self, self._io)
return
def __enter__(self, *args, **kwargs):
redirect_stdout.__enter__(self, *args, **kwargs)
return self
def __exit__(self, *args, **kwargs):
self.data += self._io.getvalue()
redirect_stdout.__exit__(self, *args, **kwargs)
return
def start(self):
self.__enter__()
return self
def stop(self):
self.__exit__(None, None, None)
return
例子:
# Using with...as
with capture() as txt1:
print('Assign these lines')
print('to a variable')
# Using start()...stop()
txt2 = capture().start()
print('This works')
print('the same way')
txt2.stop()
print('Saved in txt1:')
print(txt1)
print('Saved in txt2:')
print(txt2)
这在Sciris中实现为sc.capture()。