如何使用上下文管理器支持可选的 stdinstdout？

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 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')):
    ...

The stream_wrapper 类的效果，我觉得太过戏剧化了（假设它真的有效：我还没有测试过！）。

有没有更简单的方法来获得同样的结果？

重要的是。 任何解决这个问题的方法都必须注意永远不要关闭。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 出口时。 包裹 yield 的 open 在...中 with 如果使用，确保它是封闭的。