用于选择性重定向 I/O 的上下文管理器

Question

我经常遇到这样的情况：根据某些命令行参数，输入可能来自文件或标准输入。输出也是如此。我真的很喜欢 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)

python 标准库中有类似 DummyManager 的东西吗？实现上下文管理器协议的东西，但仅从其方法返回固定值__enter__？我想这样一个类最有可能的位置是contextlib，并且由于我在那里没有找到类似的东西，也许没有这样的东西。您还可以建议其他优雅的解决方案吗？

Answer 1

在你的情况下，你可以使用fileinputmodule：

from fileinput import FileInput

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

如果args.infile='-'那么它使用sys.stdin. 您可以传递多个文件名。如果没有文件名，它将使用命令行或标准输入中给出的文件名。inplace=Truesys.stdout

或者您可以保留文件原样：

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)