可选重定向 I/O 的上下文管理器答案

【问题标题】：Context manager for optionally redirected I/O可选重定向 I/O 的上下文管理器
【发布时间】：2014-04-25 19:10:15
【问题描述】：

我经常遇到这样的情况，即根据某些命令行参数，输入可能来自文件或标准输入。输出也是如此。我真的很喜欢 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，因为我在那里没有找到这样的东西，也许没有这样的东西。您还可以提出其他优雅的解决方案吗？

【问题讨论】：

标签： python python-3.x stdin with-statement contextmanager

【解决方案1】：

在你的情况下，你可以使用fileinput module:

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 重定向到输入文件。您可以传递多个文件名。如果没有文件名，则使用在命令行或标准输入中给出的文件名。

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

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

对于大多数可以使用标准输出来写入结果的程序来说应该没问题。

为避免关闭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)

【讨论】：

【解决方案2】：

用@contextlib.contextmanager decorator 创建一个很简单：

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)

【讨论】：

这适用于 Python 2.7 —FileInput() 直到 3.2 才获得上下文管理器支持。

【解决方案3】：

在您的情况下，您不需要虚拟经理。 sys.stdin，类似文件，可以用作上下文管理器。

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

需要注意的是，当 __exit__ing 块时，sys.stdin 会得到 closed（而您通常不需要/不想自己关闭它），但这应该不是问题。

【讨论】：

【解决方案4】：

现在存在各种ArgParse 包装器和替代品，它们很好地支持了这一点。我喜欢click：

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

如果filename 是-，这将处理sys.stdin，否则回退到常规open，其中close 隐藏在上下文管理器的finally 部分中。

【讨论】：