在Python中从脚本捕获标准输出

假设有一个脚本正在做这样的事情:

# module writer.py import sys def write(): sys.stdout.write("foobar") 

现在假设我想捕获write函数的输出并将其存储在一个variables中作进一步处理。 天真的解决scheme是:

 # module mymodule.py from writer import write out = write() print out.upper() 

但是这不起作用。 我想出了另一个解决scheme,它的工作原理,但请让我知道是否有更好的方法来解决这个问题。 谢谢

 import sys from cStringIO import StringIO # setup the environment backup = sys.stdout # #### sys.stdout = StringIO() # capture output write() out = sys.stdout.getvalue() # release output # #### sys.stdout.close() # close the stream sys.stdout = backup # restore original stdout print out.upper() # post processing 

设置stdout是一个合理的方法来做到这一点。 另一个是运行它作为另一个过程:

 import subprocess proc = subprocess.Popen(["python", "-c", "import writer; writer.write()"], stdout=subprocess.PIPE) out = proc.communicate()[0] print out.upper() 

这里是你的代码的上下文pipe理器版本。 它产生了两个值的列表; 第一个是stdout,第二个是stderr。

 import contextlib @contextlib.contextmanager def capture(): import sys from cStringIO import StringIO oldout,olderr = sys.stdout, sys.stderr try: out=[StringIO(), StringIO()] sys.stdout,sys.stderr = out yield out finally: sys.stdout,sys.stderr = oldout, olderr out[0] = out[0].getvalue() out[1] = out[1].getvalue() with capture() as out: print 'hi' 

对于未来的访问者:Python 3.4 contextlib通过redirect_stdout上下文pipe理器直接提供这个(请参阅Python的contextlib帮助 ):

 f = io.StringIO() with redirect_stdout(f): help(pow) s = f.getvalue() 

这是我原来的代码的装饰对应。

writer.py保持不变:

 import sys def write(): sys.stdout.write("foobar") 

mymodule.py sligthly得到修改:

 from writer import write as _write from decorators import capture @capture def write(): return _write() out = write() # out post processing... 

这里是装饰者:

 def capture(f): """ Decorator to capture standard output """ def captured(*args, **kwargs): import sys from cStringIO import StringIO # setup the environment backup = sys.stdout try: sys.stdout = StringIO() # capture output f(*args, **kwargs) out = sys.stdout.getvalue() # release output finally: sys.stdout.close() # close the stream sys.stdout = backup # restore original stdout return out # captured output wrapped in a string return captured 

或者,也许使用已经存在的function…

 from IPython.utils.capture import capture_output with capture_output() as c: print('some output') c() print c.stdout 

这里的问题(如何redirect输出,而不是tee部分的例子)使用os.dup2redirect操作系统级别的stream。 这很好,因为它也适用于你从程序中产生的命令。

从Python 3开始,您还可以使用sys.stdout.buffer.write()将已编码的字节string写入stdout(请参阅Python 3中的stdout )。 当你这样做时,简单的StringIO方法不起作用,因为sys.stdout.encodingsys.stdout.buffer都不可用。

从Python 2.6开始,您可以使用TextIOBase API ,其中包含缺less的属性:

 import sys from io import TextIOWrapper, BytesIO # setup the environment old_stdout = sys.stdout sys.stdout = TextIOWrapper(BytesIO(), sys.stdout.encoding) # do some writing (indirectly) write("blub") # get output sys.stdout.seek(0) # jump to the start out = sys.stdout.read() # read output # restore stdout sys.stdout.close() sys.stdout = old_stdout # do stuff with the output print(out.upper()) 

此解决scheme适用于Python 2> = 2.6和Python 3.请注意,我们的sys.stdout.write()只接受unicodestring,而sys.stdout.buffer.write()只接受字节string。 对于旧代码,情况可能不是这样,但是对于在没有变化的情况下在Python 2和3上运行的代码通常是这种情况。

如果您需要支持直接向stdout发送字节string而不使用stdout.buffer的代码,则可以使用以下变体:

 class StdoutBuffer(TextIOWrapper): def write(self, string): try: return super(StdoutBuffer, self).write(string) except TypeError: # redirect encoded byte strings directly to buffer return super(StdoutBuffer, self).buffer.write(string) 

您不必设置sys.stdout.encoding的缓冲区编码,但是当使用此方法testing/比较脚本输出时,这有帮助。

我想你应该看看这四个对象:

 from test.test_support import captured_stdout, captured_output, \ captured_stderr, captured_stdin 

例:

 from writer import write with captured_stdout() as stdout: write() print stdout.getvalue().upper() 

UPD:Eric在评论中说,不应该直接使用它,所以我复制并粘贴它。

 # Code from test.test_support: import contextlib import sys @contextlib.contextmanager def captured_output(stream_name): """Return a context manager used by captured_stdout and captured_stdin that temporarily replaces the sys stream *stream_name* with a StringIO.""" import StringIO orig_stdout = getattr(sys, stream_name) setattr(sys, stream_name, StringIO.StringIO()) try: yield getattr(sys, stream_name) finally: setattr(sys, stream_name, orig_stdout) def captured_stdout(): """Capture the output of sys.stdout: with captured_stdout() as s: print "hello" self.assertEqual(s.getvalue(), "hello") """ return captured_output("stdout") def captured_stderr(): return captured_output("stderr") def captured_stdin(): return captured_output("stdin") 

我喜欢上下文pipe理解决scheme,但是如果你需要使用打开的文件和fileno支持存储缓冲区,你可以做这样的事情。

 import six from six.moves import StringIO class FileWriteStore(object): def __init__(self, file_): self.__file__ = file_ self.__buff__ = StringIO() def __getattribute__(self, name): if name in { "write", "writelines", "get_file_value", "__file__", "__buff__"}: return super(FileWriteStore, self).__getattribute__(name) return self.__file__.__getattribute__(name) def write(self, text): if isinstance(text, six.string_types): try: self.__buff__.write(text) except: pass self.__file__.write(text) def writelines(self, lines): try: self.__buff__.writelines(lines) except: pass self.__file__.writelines(lines) def get_file_value(self): return self.__buff__.getvalue() 

使用

 import sys sys.stdout = FileWriteStore(sys.stdout) print "test" buffer = sys.stdout.get_file_value() # you don't want to print the buffer while still storing # else it will double in size every print sys.stdout = sys.stdout.__file__ print buffer