Pythonでpasteコマンドのように標準入力または可変個のファイルを同時に読み込む方法

Pythonでファイルをopenするときはwith文を使うのが良い習慣とされている。しかし、この構文で可変個のファイルを同時にopenするのは一筋縄ではいかない。このように書いてももちろん動かない。

with [open(fn) for fn in file_names] as files:

例えばpasteコマンドのような処理を行うとき、同時openはどうしても必要になってくる。そこで本記事では、Pythonでpasteコマンドのような動作をwith文を使って実現する方法について書く。

※ ファイルを1個ずつopenしたいときはfileopenモジュールを使えばよい（参考: Pythonでcatみたいに標準入力またはファイル名指定でテキストを読み込む方法）。
※※ ファイルを定数個だけopenするならwith open(a) as f1, open(b) as f2:などと書けばよい

contextlibモジュールを使え！

paste.py

import argparse
import contextlib


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('fs', nargs='*', type=argparse.FileType(),
                        default='-')
    parser.add_argument('--delimiter', default='\t')
    args = parser.parse_args()

    with contextlib.ExitStack() as stack:
        for f in args.fs:
            stack.enter_context(f)
        paste(args.fs, args.delimiter)


if __name__ == "__main__":
    main()

add_argument()でtype=argparse.FileType(), default='-'を指定したときの動作はこちらの記事にある通りである。そして今回、nargs='*'と組み合わせることで、可変個のコマンドライン引数からファイルオブジェクトのリストが得られる。

そして、最も伝えたいのがこのwith文の部分である。そもそもwith文とは、ブロックの実行を「コンテキストマネージャ」のメソッドでラップするために使われるものである。そしてこのコンテキストマネージャ型を簡単に実装するためのモジュールがcontextlibである。

今回のように可変個のファイルをwith文で扱うには、まずExitStack()でコンテキストマネージャを生成し、enter_context()を使って管理したいオブジェクトを追加していけばよい。これでwithブロックを抜けるときにしっかりcloseが呼ばれる。

pasteの実装

本記事の趣旨からそれるが、paste()の実装例はこちら。

from itertools import zip_longest


def rstrip(x):
    if x is None:
        return ''
    else:
        return x.rstrip()


def paste(files, delimiter):
    for lines in zip_longest(*files):
        stripped = map(rstrip, lines)
        print(delimiter.join(stripped))

余談

言語処理100本ノックの問13もこのコードでOK