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

Question

假设我想实现一个具有以下签名的 Python 脚本：

myscript.py INPUT OUTPUT

...其中INPUT和分别OUTPUT代表脚本将读取和写入的文件的路径。

用于实现具有此类签名的脚本的代码可能具有以下构造：

with open(inputarg, 'r') as instream, open(outputarg, 'w') as outstream:
    ...

...这里的inputarg和变量保存通过其和命令行参数outputarg传递给脚本的文件路径（字符串） 。INPUTOUTPUT

---

到目前为止没有什么特别或不寻常的事情。

但现在，假设对于脚本的版本 2，我想为用户提供-为其一个（或两个）参数传递特殊值的选项，以指示脚本应分别读取stdin和写入stdout。

换句话说，我希望以下所有形式都能产生相同的结果：

myscript.py INPUT OUTPUT
myscript.py   -   OUTPUT  <INPUT
myscript.py INPUT   -             >OUTPUT
myscript.py   -     -     <INPUT  >OUTPUT

现在，with之前的说法已经不再适用。一方面，表达式open('-', 'r')oropen('-', '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')):
    ...

这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在opena 中with可确保其在使用时关闭。