捕获标准输出时python2和python3之间的StringIO可移植性

Question

我写了一个python软件包，设法使它与python 2.7和python 3.4完全兼容，但到目前为止，有一个例外使我感到困惑。该软件包包括一个命令行脚本，在我的单元测试中，我使用此代码运行脚本的主例程，同时覆盖sys.argv以传递argparse的命令行参数，并捕获脚本的stdout进行比较：

@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之间移植。

顺便说一句，我必须承认我完全不太了解装饰，上下文管理器或“ yield”关键字。我的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下已部分中断。有人可以帮助我找出如何以可移植的Python方式解决此stdout重定向问题的方法，最好不要添加任何其他外部依赖项吗？

Answer 1

您将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对象与Python 2等价io.BytesIO；它要求您编写纯字节串，而不是unicode对象。

您也可以io.BytesIO在Python 2中使用，但是随后您想测试if是否sys.stdout是io.TextIOBase子类；如果不是，则用binary 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()