Pythonでcatみたいに標準入力またはファイル名指定でテキストを読み込む方法

catなどのコマンドは、ファイル名を引数にとって読み込みを行いますが、ファイル名が与えられない場合は標準入力を読む、という挙動をします。今回はこれをPythonで実現する方法について調べました。

$ cat input.txt
$ cat < input.txt
$ python hoge.py input.txt
$ python hoge.py < input.txt

0. 基本

Pythonでは標準入力もテキストファイルへのポインタもfile object（細かく言うと_io.TextIOWrapperクラスのインスタンス）なので、次のような関数は標準入力sys.stdinが引数に与えられても動作します。

def process(fp):
    for line in fp:
        print(line, end='')

本記事はこの外側の処理をどうするか？というお話で、3通りの方法を検討していきます。

1. 素朴に場合分け

hoge.py

import argparse
import sys


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('filename', nargs='?')
    args = parser.parse_args()

    if args.filename is None:
        process(sys.stdin)
    else:
        with open(args.filename) as f:
            process(f)

if __name__ == "__main__":
    main()

コマンドライン引数が与えられたかどうかで分岐する素直な方法です。なお本稿では、open()やprocess()実行時のOSErrorに対する特別な例外処理は考えないものとしますが、process()で例外が発生したときでもファイルをcloseするのはマストとします。
ここでargparseについて補足しておくと、もしnargs='?'を省略するとfilenameが必須になってしまいます。

この方法はとってもわかりやすいのですが、釈然としないのは、process()を2箇所に記述しているところです。もう少し工夫すれば1箇所にまとめられますが、それでもif-elseがやや冗長に感じます。

2. argparse.Filetype()

hoge.py

import argparse
import sys


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('infile', nargs='?', type=argparse.FileType(),
                        default=sys.stdin)
    args = parser.parse_args()
    with args.infile as f:
        process(f)

少々見慣れない形ですが、全体としては非常にスッキリしました。

type=argparse.FileType()とすることで、parse_args()の瞬間にopen()が呼び出され、その戻り値が格納されます。このとき、与えられたファイル名が-のときは標準入出力に変換されます。そのためdefault='-'としてしまえば記述がちょっと減ってimport sysも不要になるんですが、知らない人にとってはわかりにくいですし、ドキュメントの例でもdefault=sys.stdinと明示されています。

そしてwithの直後はopen()を書くのが普通ですが、ファイルオブジェクトそのものにしても戻り値は同じなので、同じ動作になります。asを使う意義は薄れますが。

sys.stdinは通常わざわざwith文と共に使ったりclose()を呼び出したりしませんが、TextIOWrapperなのでちゃんと定義されています。おそらくこういうときに使うものなのでしょう。（要出典）

3. argparse + fileinput

レアケースだと思いますが、ファイル名を複数指定できるようにし、それらを1ファイルずつ読み込みたい、また1個も指定されなかったら標準入力を使うようにしたい、という場合があります。そういう場合はfileinputモジュールを併用すると簡単です。

hoge.py

import argparse
import fileinput


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('filenames', nargs='*')
    args = parser.parse_args()
    with fileinput.input(args.filenames) as f:
        process(f)

fileinput.input()は列挙されているファイル名を次々とopenし、もしファイル名が与えられてないときは標準入力に変換してくれます。

ファイル名として-が指定された場合も標準入力が呼ばれるので、$ python hoge.py input.txt -などと実行するとinput.txtの読み込み後に標準入力も読み込むことになります。まぁこれは大して問題にはならないでしょう！

個人的に難点だと思うのは、このfileinputモジュールは知らないと標準入力に対応しているように見えないところです。これそんなに有名なモジュールじゃないと思うんですよね…。

まとめ

1個のファイルか標準入力を読むときは2の方法、複数のファイルまたは標準入力を読むときは3の方法を使うのが良いと思います。

(11/23 追記) 複数のファイルを同時に読み込みたい場合はこちら→Pythonでpasteコマンドのように標準入力または可変個のファイルを同時に読み込む方法

疑問

同じことを他の言語でやろうとすると、大抵は1のようにif-elseでやるしかなく、結構大変なのでは？Perlは簡単らしいですが。