从文件或 STDIN 读取

Question

我编写了一个命令行实用程序，它使用 getopt 来解析命令行上给出的参数。我还希望文件名是一个可选参数，例如它在其他实用程序中，如 grep、cut 等。所以，我希望它具有以下用法

tool -d character -f integer [filename]

如何实现以下功能？

• 如果给出了文件名，则从文件中读取。
• 如果没有给出文件名，则从 STDIN 读取。

Answer 1

fileinput 模块可以做你想做的事——假设非选项参数在args 中，那么：

import fileinput
for line in fileinput.input(args):
    print line

如果args 为空，则fileinput.input() 将从标准输入读取；否则它会依次从每个文件中读取，类似于 Perl 的while(<>)。

Answer 2

用最简单的话来说：

import sys
# parse command line
if file_name_given:
    inf = open(file_name_given)
else:
    inf = sys.stdin

此时您将使用inf 从文件中读取。根据是否给出了文件名，这将从给定文件或标准输入中读取。

当你需要关闭文件时，你可以这样做：

if inf is not sys.stdin:
    inf.close()

但是，在大多数情况下，如果您完成了关闭 sys.stdin，它是无害的。

Answer 3

我喜欢使用上下文管理器的一般习惯用法，但是当您退出 with 语句时，（太）琐碎的解决方案最终会关闭 sys.stdin，我想避免这种情况。

借鉴this answer，这里有一个解决方法：

import sys
import contextlib

@contextlib.contextmanager
def _smart_open(filename, mode='Ur'):
    if filename == '-':
        if mode is None or mode == '' or 'r' in mode:
            fh = sys.stdin
        else:
            fh = sys.stdout
    else:
        fh = open(filename, mode)
    try:
        yield fh
    finally:
        if filename != '-':
            fh.close()
    
if __name__ == '__main__':
    args = sys.argv[1:]
    if args == []:
        args = ['-']
    for filearg in args:
        with _smart_open(filearg) as handle:
            do_stuff(handle)

我想你可以实现something similar with os.dup()，但我编写的代码变得更加复杂和神奇，而上面的代码有些笨拙但非常简单。

Answer 4

我更喜欢使用“-”作为你应该从标准输入读取的指示符，它更明确：

import sys
with open(sys.argv[1], 'r') if sys.argv[1] is not "-" else sys.stdin as f:
    pass # do something here

Answer 5

要利用python的with语句，可以使用以下代码：

import sys
with open(sys.argv[1], 'r') if len(sys.argv) > 1 else sys.stdin as f:
    # read data using f
    # ......

Answer 6

不是直接的答案，而是相关的。

通常，当您编写 python 脚本时，您可以使用 argparse 包。 如果是这种情况，您可以使用：

parser = argparse.ArgumentParser()
parser.add_argument('infile', nargs='?', type=argparse.FileType('r'), default=sys.stdin)

'?'。如果可能，将从命令行使用一个参数， 并作为单个项目生产。如果不存在命令行参数， 将产生默认值。

这里我们将默认设置为sys.stdin;

所以如果有文件，它将读取它，如果没有，它将从标准输入获取输入“注意：我们在上面的示例中使用了位置参数”

更多访问：https://docs.python.org/2/library/argparse.html#nargs

Answer 7

切换到argparse（它也是标准库的一部分）并使用 argparse.FileType，默认值为标准输入：

import  argparse, sys

p = argparse.ArgumentParser()
p.add_argument('input', nargs='?',
  type=argparse.FileType(), default=sys.stdin)
args = p.parse_args()

print(args.input.readlines())

这不会让你为标准输入指定编码和其他参数， 然而;如果你想这样做，你需要使参数非可选 并让FileType 在- 作为 论据：

p.add_argument('input', type=FileType(encoding='UTF-8'))

注意后一种情况不会支持二进制模式 ('b') I/O。如果 你只需要，你可以使用上面的默认参数技术，但是 提取二进制 I/O 对象，例如，default=sys.stdout.buffer for 标准输出。但是，如果用户无论如何指定-，这仍然会中断。 （- stdin/stdout 总是包裹在 TextIOWrapper 中。）

如果您希望它与 - 一起使用，或者您需要任何其他参数 在打开文件时提供，如果它被包装，您可以修复参数 错误：

p.add_argument('output', type=argparse.FileType('wb'))
args = p.parse_args()
if hasattr(args.output, 'buffer'):
    #   If the argument was '-', FileType('wb') ignores the 'b' when
    #   wrapping stdout. Fix that by grabbing the underlying binary writer.
    args.output = args.output.buffer

（帽子提示medhat 提到add_argument()的type参数。）

Answer 8

类似：

if input_from_file:
    f = open(file_name, "rt")
else:
    f = sys.stdin
inL = f.readline()
while inL:
    print inL.rstrip()
    inL = f.readline()