StringIO portability between python2 and python3 when capturing stdout

Question

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:

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

Incidentally, my complete package where this code comes from is here:

https://github.com/NF6X/pyImageDisk

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?

Answer 1

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()