Как поддерживать опциональный stdin / stdout с помощью контекстных менеджеров? - PullRequest
Новая
       95

Как поддерживать опциональный stdin / stdout с помощью контекстных менеджеров?

1 голос
/ 24 апреля 2020

Предположим, я хочу реализовать скрипт Python со следующей подписью: 

myscript.py INPUT OUTPUT

... где INPUT и OUTPUT обозначают пути к файлам сценарий будет читать и записывать соответственно.

Код для реализации сценария с такой подписью может содержать следующую конструкцию: 

with open(inputarg, 'r') as instream, open(outputarg, 'w') as outstream:
    ...

.. . Здесь переменные inputarg и outputarg содержат пути к файлам (которые являются строками), переданные сценарию через аргументы командной строки INPUT и OUTPUT.

Ничего особенного или пока что необычно.

Но теперь предположим, что для версии 2 скрипта я хочу дать пользователю возможность передать специальное значение - для любого (или обоих) своих аргументов, чтобы указать, что скрипт должен, соответственно, читать из stdin и записывать в stdout.

Другими словами, я хочу, чтобы все приведенные ниже формы давали одинаковые результаты: 

myscript.py INPUT OUTPUT
myscript.py   -   OUTPUT  <INPUT
myscript.py INPUT   -             >OUTPUT
myscript.py   -     -     <INPUT  >OUTPUT

Теперь приведенное ранее выражение with больше не подходит. С одной стороны, выражение open('-', 'r') или open('-', 'w') может вызвать исключение: 

FileNotFoundError: [Errno 2] No such file or directory: '-'

Я не смог придумать удобный способ расширения приведенная выше конструкция на основе with для размещения желаемой новой функциональности.

Например, этот вариант не будет работать (помимо громоздкости), поскольку sys.stdin и sys.stdout не реализуют интерфейс диспетчера контекста: 

with sys.stdin if inputarg == '-' else open(inputarg, 'r'), \
        sys.stdout if outputarg == '-' else open(outputarg, 'w'):
    ...

Единственное, что я могу придумать (возможно), это определить минимальный сквозной класс-оболочку, который реализует интерфейс диспетчера контекста, например: 

class stream_wrapper(object):

    def __init__(self, stream):
        self.__dict__['_stream'] = stream

    def __getattr__(self, attr):
        return getattr(self._stream, attr)

    def __setattr__(self, attr, value):
        return setattr(self._stream, attr, value)

    def close(self, _std=set(sys.stdin, sys.stdout)):
        if not self._stream in _std:
            self._stream.close()

    def __enter__(self):
        return self._stream

    def __exit__(self, *args):
        return self.close()

... и затем напишите оператор with следующим образом: 

with stream_wrapper(sys.stdin if inputarg == '-' else open(inputarg, 'r')), \
        stream_wrapper(sys.stdout if outputarg == '-' else open(outputarg, 'w')):
    ...

Класс stream_wrapper вызывает у меня много драмы за то, что он достигает (при условии, что он работает вообще: я не проверял это!).

Есть ли более простой способ получить те же результаты?

ВАЖНО: Любое решение этой проблемы должно позаботиться никогда не закрывать sys.stdin или sys.stdout.

1 Ответ

2 голосов
/ 24 апреля 2020

Используя contextlib.contextmanager , этим можно управлять с помощью чего-то вроде: 

from contextlib import contextmanager
import sys

@contextmanager
def stream(arg,mode='r'):
    if mode not in ('r','w'):
        raise ValueError('mode not "r" or "w"')
    if arg == '-':
        yield sys.stdin if mode == 'r' else sys.stdout
    else:
        with open(arg,mode) as f:
            yield f

with stream(sys.argv[1],'r') as fin,stream(sys.argv[2],'w') as fout:
        for line in fin:
            fout.write(line)

Если вы не знакомы с contextmanager, он в основном запускает код до yield на вход и после yield на выходе. Обертывание yield open в with гарантирует, что оно закрыто, если используется.

...