Диспетчер контекста для опционально перенаправленного ввода-вывода

Я часто сталкиваюсь с ситуацией, когда в зависимости от какого-либо аргумента командной строки ввод может происходить из файла или из стандартного ввода. То же самое касается вывода. Мне очень нравится, как работают менеджеры контекста в Python 3, и поэтому стараюсь сделать все мои open называет частью некоторых with заявление. Но в этом случае у меня проблемы.

if args.infile:
    with open(args.infile, "r") as f:
        process(f)
else:
    process(sys.stdin)

уже неуклюжий, и с входом и выходом мне придется обслуживать четыре комбинации. Я хотел бы что-то проще, например

with (open(args.infile, "r") if args.infile
      else DummyManager(sys.stdin)) as f:
    process(f)

Есть ли что-то вроде этого DummyManager в стандартных библиотеках python? Что-то, что реализует протокол менеджера контекста, но только для возврата фиксированного значения из его __enter__ метод? Я предполагаю, что наиболее вероятным местоположением для такого класса будет contextlib, и, поскольку я не нашел ничего подобного, возможно, такого не существует. Есть ли другие элегантные решения, которые вы можете предложить?

4 ответа

Решение

В вашем случае вы могли бы использовать fileinput модуль:

from fileinput import FileInput

with FileInput(args.infile) as file:
    process(file)
# sys.stdin is still open here

Если args.infile='-' тогда он использует sys.stdin, Если inplace=True параметр, то он перенаправляет sys.stdout во входной файл. Вы можете передать несколько имен файлов. Если нет имен файлов, он использует имена файлов, заданные в командной строке или в stdin.

Или вы можете оставить файл как есть:

import sys
import argparse

parser = argparse.ArgumentParser()
parser.add_argument('--log', default=sys.stdout, type=argparse.FileType('w'))
args = parser.parse_args()
with args.log:
    args.log.write('log message')
# sys.stdout may be closed here

Это должно быть хорошо для большинства программ, где stdout может использоваться для записи результата.

Чтобы избежать закрытия sys.stdin / sys.stdout Вы могли бы использовать ExitStack включить контекстные менеджеры условно:

from contextlib import ExitStack

with ExitStack() as stack:
    if not args.files:
       files = [sys.stdin]
    else:
       files = [stack.enter_context(open(name)) for name in args.files]

    if not args.output:
       output_file = sys.stdout
       stack.callback(output_file.flush) # flush instead of closing 
    else:
       output_file = stack.enter_context(open(args.output, 'w'))

    process(files, output_file)

Это тривиально, чтобы создать один с @contextlib.contextmanager декоратор:

from contextlib import contextmanager

@contextmanager
def dummy_manager(ob):
    yield ob

Это оно; это создает менеджер контекста, который ничего не делает, но вручает вам ob назад, а __exit__ Обработчик точно ничего не делает.

Я бы использовал это так:

f = open(args.infile, "r") if args.infile else dummy_manager(sys.stdin)
with f:
    process(f)

Вам не нужен фиктивный менеджер в вашем случае. sys.stdinбудучи подобным файлу, может использоваться как менеджер контекста.

with (open(args.infile, "r") if args.infile else sys.stdin) as f:
    process(f)

Стоит отметить, что когда __exit__В блоке sys.stdin получает closed (хотя обычно вам не нужно / не хотите закрывать его самостоятельно), но это не должно быть проблемой.

Различный ArgParse В настоящее время существуют обертки и замены, которые прекрасно это поддерживают. мне нравится click:

with click.open_file(filename) as lines:
    for line in lines:
        process(line)

Это справится sys.stdin если filename является - а в противном случае отступить к обычному open с close скрыто в контексте менеджера finally часть.

Другие вопросы по тегам