I have written a python package which I have managed to make fully compatible with both python 2.7 and python 3.4, with one exception that is stumping me so far. The package includes a command line script, and in my unit tests I use this code to run the script's main routine while overriding sys.argv to pass command line arguments for argparse, and capturing the script's stdout for comparison:
@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...
This works well in python 3.4, but in python 2.7 it generates an error:
TypeError: unicode argument expected, got 'str'
I haven't managed to figure out a way to capture stdout from arbitrary functions which is portable between python 2.7 and python 3.4.
As an aside, I have to admit that I don't understand decorations, context managers or the "yield" keyword very well at all. The inspiration for my runmain() function came from:
Incidentally, my complete package where this code comes from is here:
At the moment, its unit tests are partially broken under python 2.7 because of this issue. Can anybody help me figure out how to solve this stdout redirection problem in a portable, pythonic manner, preferably without adding any more external dependencies?
解决方案
You replaced the Python 2 bytes-only sys.stdout with one that only takes Unicode. You'll have to adjust your strategy on the Python version here, and use a different object:
try:
# Python 2
from cStringIO import StringIO
except ImportError:
# Python 3
from io import StringIO
and remove the io. prefix in your context manager:
origout, sys.stdout = sys.stdout, StringIO()
The cStringIO.StringIO object is the Python 2 equivalent of io.BytesIO; it requires that you write plain bytestrings, not aunicode objects.
You can also use io.BytesIO in Python 2, but then you want to test if sys.stdout is a io.TextIOBase subclass; if it is not, replace the object with a binary BytesIO, object, otherwise use a StringIO object:
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()