在捕获标准输出时，StringIO在Python2和Python3之间的可移植性问题

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'

我还没有找到一种可在Python 2.7和Python 3.4之间通用的从任意函数中捕获标准输出(stdout)的方法。

顺便提一下，我必须承认，我不太理解装饰器(decorations)，上下文管理器(context managers)或“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惯例的方式解决这个标准输出重定向问题吗？最好是不添加任何外部依赖项。

- NF6X

Answer 1

您用只支持Unicode的一个替代了Python 2中只支持字节的sys.stdout。因此，您需要调整在这里使用的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对象。

在Python 2中也可以使用io.BytesIO，但是您需要测试sys.stdout是否为 io.TextIOBase子类; 如果不是，请将对象替换为一个二进制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()

- Martijn Pieters

---

在这里稍微提一下，我相信可以放弃 origout 并使用默认备份 sys.__stdout__，它也存在于两种 Python 中。 - Dimitris Fasarakis Hilliard

谢谢，问题已经解决了！接下来，我发现Python 2.7中不包含unittest.assertRegex()函数，而unittest.assertRegexpMatches()在3.4中会生成一个弃用警告。离成功更进一步了... - NF6X

@Jim 不行，因为那会忽略之前的标准输出捕获。 - Martijn Pieters

Answer 2

你试过了吗？（可以在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'))

- mementum

---

我根据您的建议尝试了，但在2.7下它并没有改变行为。不管怎样，非常感谢您的快速回复！ - NF6X

问题在于，在Python 2中，sys.stdout 必须 是一个接受 str（字节）而不是 unicode 的对象。 - Martijn Pieters