如何使用上下文管理器支持可选的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 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之后。在with中包装open的yield可以确保它在使用时是关闭的