Python: 適宜ファイルを差し替えながら大量のデータを書き込む

概要

大きなデータを書き込むときに、1つのファイルに書き込んでしまうと、読み込みのときに処理時間が膨大になったり、あるいはメモリのエラーで実行がそもそも不能になる可能性があります。

そこで、書き出し量がある一定量を超えたときに、出力ファイルを切り替えて、1ファイルが大きくならないようにするロジックを書くことがあるかもしれません。普通は主要処理の中に直接ロジックを書き込んだりすることも多いと思いますが、クラス化して使いやすくしました。

ファイル複数書き出しのクラス

コード

class FileWriter:

    def __init__(self, output_directory, path_format, threshold):
        self.output_directory = output_directory
        self.path_format = path_format
        self.threshold = threshold
        self.count = 0
        self.total_count = 0
        self.file_count = 0
        self.file_object = None

    def write(self, string, count = 1):
        if self.file_object is None:
            self.file_count += 1
            self.file_object = open(f"{self.output_directory}/{self.path_format.format(self.file_count)}", "w")

        self.file_object.write(string)
        self.count += count
        self.total_count += count

        if self.count >= self.threshold:
            self.file_object.close()
            self.file_object = None
            self.count = 0

    def finish(self):
        if self.file_object is not None:
            self.file_object.close()
            self.file_object = None

    def __enter__(self):
        return self

    def __exit__(self, ex_type, ex_value, trace):
        self.finish()

使い方

with の中で使います。書き込み用のディレクトリ (output_directory) を指定して、あとはひたすら write を実行していくだけです。 path_format に従ったファイル名で出力がされます。 write するたびに数値を与え、その合計が threshold を超えたら、ファイルの番号を1増やして新たに書き込みを始めていきます。

使ってみる

0 ~ 1023 までの数字を1行ずつ書いていきます。write に count の値を与えていないため、デフォルト値の1がカウント加算されます。カウントが threshold に指定している 100を超えたところで新しいファイルに移行し、書き込みを続けます。

with FileWriter("out", "single_{:04d}.txt", 100) as fw:
    for i in range(1024):
        fw.write(str(i) + "\n")

これにより、 out/single_0001.txt ~ out/single_0011.txt というファイルに、数字が連番で書き出されます。

さて、1つの文章であるため、ファイルが分割されないように出力したいが、複数行の文字列であるため、行数分だけカウントは増やしたいということがあるかもしれません。そういう場合は len などを駆使しつつ、 write の引数 count に値を指定すればオーケーです。

with FileWriter("out", "multiple_{:04d}.txt", 100) as fw:
    values = [str(i) for i in range(40)]
    string = "\n".join(values) + "\n"
    for i in range(10):
        fw.write(string, len(values))

0 ~ 39 までの数字を40行でまとめて書き出すため、100を超えるタイミング、つまり 40 * 3 = 120 のときにファイルが切り替わります。 out/multiple_0001.txt ~ out/multiple_0004.txt が出来上がります。