如何使用上下文管理器支持可选 stdin/stdout？

Question

假设我想实现一个具有以下签名的 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，如下所示：

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 class 给我留下了深刻的印象，因为它所取得的成就非常戏剧化（假设它完全有效：我还没有测试过！）。

是否有更简单的方法来获得相同的结果？

重要提示： 此问题的任何解决方案都必须注意永远不要关闭 sys.stdin 或 sys.stdout.

Answer 1

使用 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 之后。将 open 的 yield 包装在 with 中确保它在使用时被关闭。

如何使用上下文管理器支持可选 stdin/stdout？

How to support optional stdin/stdout using context managers?

python

stdin

stdout

contextmanager