可选重定向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

现在有很多不同的 ArgParse 的包装和替代方案，它们都能很好地支持这个功能。我个人比较喜欢 click：

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

这个代码会处理 sys.stdin，如果 filename 是 -，否则就会使用普通的 open 方法，并且在上下文管理器的 finally 部分隐藏了 close 的操作。

Answer 2

在你的情况中，其实不需要一个虚假的管理器。因为 sys.stdin 是一种类似文件的东西，可以直接用作上下文管理器。

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

需要注意的是，当你退出这个代码块时，sys.stdin 会被自动关闭（虽然通常你不需要自己去关闭它），但这应该不会造成什么问题。

Answer 3

用@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)

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

或者你也可以让文件保持不变：

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)