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

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


person NF6X    schedule 19.01.2016    source источник


Ответы (2)


Вы заменили 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 является эквивалентом io.BytesIO в Python 2; это требует, чтобы вы писали простые строки байтов, а не объекты aunicode.

Вы также можете использовать io.BytesIO в Python 2, но тогда вы хотите проверить, является ли sys.stdout _ 10_ подкласс; если это не так, замените объект двоичным 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()
person Martijn Pieters    schedule 19.01.2016
comment
Немного в стороне, я почти уверен, что можно было бы отказаться от origout и использовать резервную копию sys.__stdout__ по умолчанию, которая также присутствует в обоих Pythons. - person Dimitris Fasarakis Hilliard; 19.01.2016
comment
Спасибо, что исправили! Затем я обнаружил, что python 2.7 не включает unittest.assertRegex (), а unittest.assertRegexpMatches () генерирует предупреждение об устаревании в 3.4. Приближаемся ... - person NF6X; 19.01.2016
comment
@Jim нет, потому что это игнорировало любой предыдущий захват stdout. - person Martijn Pieters; 19.01.2016

Вы пробовали? (Можно оставить в вашем коде под 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'))
person mementum    schedule 19.01.2016
comment
Я попробовал это на основе вашего предложения, но это не повлияло на поведение в версии 2.7. Но все равно спасибо за очень быстрый ответ! - person NF6X; 19.01.2016
comment
Проблема в том, что в Python 2 sys.stdout должен быть объектом, который принимает str (байты), а не unicode. - person Martijn Pieters; 19.01.2016