捕获stdout

时间:2016-01-19 08:17:38

标签: python python-2.7 stdout python-3.4 stringio

我已经编写了一个python包,我已经设法与python 2.7和python 3.4完全兼容,但有一个例外是我到目前为止所困扰的。该软件包包含一个命令行脚本,在我的单元测试中,我使用此代码运行脚本的主例程,同时重写sys.argv以传递argparse的命令行参数,并捕获脚本的标准输出比较:

@contextlib.contextmanager
def runmain(mainfunction, arglist):
    """Run mainfunction with arglist in sys.srgv, and capture stdout."""

    origargv, sys.argv   = sys.argv,   arglist
    origout,  sys.stdout = sys.stdout, io.StringIO()

    rtn = mainfunction()

    sys.stdout.seek(0)
    yield (rtn, sys.stdout.read())

    sys.stdout = origout
    sys.argv   = origargv

class test_imdutil_main(unittest.TestCase):

    def test_help(self):
        """Test -h option."""

        with runmain(imdutil_main, ['imdutil.py', '-h']) as (rtn, capture):
            # do stuff with rtn and capture...

这在python 3.4中运行良好,但在python 2.7中会产生错误:

TypeError: unicode argument expected, got 'str'

我还没有设法找到一种从任意函数中捕获stdout的方法,这些函数可以在python 2.7和python 3.4之间移植。

顺便说一句,我不得不承认,我不了解装饰品,背景管理者或者"收益"关键字非常好。我的runmain()函数的灵感来自:

http://schinckel.net/2013/04/15/capture-and-test-sys.stdout-sys.stderr-in-unittest.testcase/

顺便提一下,我的完整包来自此代码:

https://github.com/NF6X/pyImageDisk

目前,由于此问题,其单元测试在python 2.7下部分被破坏。任何人都可以帮我弄清楚如何以便携式,pythonic方式解决这个stdout重定向问题,最好不再添加任何外部依赖项?

2 个答案:

答案 0 :(得分:5)

您将仅使用Python的2个字节sys.stdout替换为仅使用Unicode。您必须在此处调整Python版本的策略,并使用其他对象:

try:
    # Python 2
    from cStringIO import StringIO
except ImportError:
    # Python 3
    from io import StringIO

并删除上下文管理器中的io.前缀:

origout,  sys.stdout = sys.stdout, StringIO()

cStringIO.StringIO对象是等效于io.BytesIO的Python 2;它要求你写简单的字节串,而不是unicode个对象。

你也可以在Python 2中使用io.BytesIO,但是你想测试sys.stdoutio.TextIOBase subclass;如果不是,请用二进制BytesIO对象替换该对象,否则使用StringIO对象:

import io

if isinstance(sys.stdout, io.TextIOBase):
    # Python 3
    origout, sys.stdout = sys.stdout, io.StringIO()
else:
    # Python 2 or an unorthodox binary stdout setup
    origout, sys.stdout = sys.stdout, io.BytesIO()

答案 1 :(得分:-1)

你试过吗? (可以留在Python 3.x下的代码中)

from __future__ import unicode_literals

使用io.StringIO时,我的代码中包含的内容使其兼容:

f = io.StringIO(datafile.read().decode('utf-8'), newline=None)

然后查看您的代码:

yield (rtn, sys.stdout.read())

可以更改为:

yield (rtn, sys.stdout.read().decode('utf-8'))