python - 相当于./foo.py <bar.png的Pythonic

我有一个从sys.stdin读取的Python程序，因此可以使用./foo.py < bar.png进行调用。如何从另一个Python模块中测试此代码？也就是说，如何在运行测试脚本时将stdin设置为指向文件的内容？我不想做类似./test.py < test.png的事情。我认为我不能使用fileinput，因为输入是二进制的，我只想处理一个文件。使用PIL中的Image.open(sys.stdin)打开文件。

最佳答案

您应该对脚本进行一般化，以便除了用作独立程序之外，还可以从测试脚本中调用它。这是执行此操作的示例脚本：

#! /usr/bin/python

import sys

def read_input_from(file):
    print file.read(),

if __name__ == "__main__":
    if len(sys.argv) > 1:
        # filename supplied, so read input from that
        filename = sys.argv[1]
        file = open(filename)
    else:
        # no filename supplied, so read from stdin
        file = sys.stdin
    read_input_from(file)

如果使用文件名调用该文件，则将显示该文件的内容。否则，将显示从标准输入读取的输入。 （对于foo.py脚本，能够在命令行上传递文件名可能是有用的改进。）

现在，在测试脚本中，您可以使用文件调用foo.py中的函数，例如：

#! /usr/bin/python

import foo

file = open("testfile", "rb")
foo.read_input_from(file)