Переносимость StringIO между python2 и python3 при захвате stdout

Question

Я написал пакет 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' 

мне не удалось выяснить способ захвата стандартного вывода произвольных функций, переносим между 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 из-за этой проблемы. Может ли кто-нибудь помочь мне разобраться, как решить эту проблему перенаправления 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.

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