Pythonのジェネレータを使って大容量ファイルを分割する

過去に下記のような、ファイルを複数に分割する方法について記事を書いた：

blog.honjala.net

この記事に書いたコードをもう少しエレガントに書けないものか、

と方法を調べていて、

Pythonのジェネレータが使えそうだったので色々試してみた。

Pythonのジェネレータについて

ジェネレータを意識して使ったことがなかったので、まずは使ってみる。

下記のページを参照した：

postd.cc

ジェネレータを使用する際は、

「ジェネレータ関数」と

ジェネレータ関数が返す「イテレータオブジェクト」の

2つの要素を使う。

yield文で返り値を返す関数を定義すると、ジェネレータ関数になる。

ジェネレータ関数は呼び出しても中の処理が実行されるわけではなく、

代わりにイテレータオブジェクトを返す。

イテレータオブジェクトをforループなりnext関数なりで使用すると、

ジェネレータ関数で定義したyield文まで処理が進み、指定した値が返される。

そこでジェネレータ関数内の処理は一旦停止し、呼び出し元にフローが戻ってくる。

再度イテレータオブジェクトを呼び出すと、処理が再開し、

次のyield文に到達するか関数内の処理が終了するまで進む。

...

分かりやすく説明できなかったので、下記のWikipediaの説明をご参照ください：

ジェネレータ (プログラミング) - Wikipedia

とりあえずジェネレータ関数を使う時のポイントは

「関数内の処理を呼び出しごとに一時停止できる」

という点で、これを活用することで大きなデータを

メモリを圧迫せずに読み込んだり処理したりできるようだ。

関数内の処理が一時停止することを確認するために、

上記の参照サイトのサンプルに手を加えて実行してみた：

def generate_nums():
    for num in range(15):
        yield num 

nums = generate_nums()

for x in nums:
    print("first loop: %d" % x)

    if x > 9:
        break

for x in nums:
    print("second loop: %d" % x)

    if x > 19: 
        break

generate_numsがジェネレータ関数で、

numsが生成されたイテレータオブジェクト。

numsを呼び出すごとにループを1回回して連番の数値を1つ返すので、

呼び出される度に0から14までの連番を返すはず。

実験として、呼び出し元は2回のループでnumsを呼び出すようにして、

返り値が10になるまで最初のループで回し、

続きを2番目のループで再開できていることを確認できるようにした。

出力は下記のようになった：

first loop: 0
first loop: 1
first loop: 2
first loop: 3
first loop: 4
first loop: 5
first loop: 6
first loop: 7
first loop: 8
first loop: 9
first loop: 10
second loop: 11
second loop: 12
second loop: 13
second loop: 14

ちゃんと再開出来てるようだ。

ファイル内容を1行ずつ返すジェネレータ

ジェネレータについて学んだところで、

実際にファイルを1行ずつ返すジェネレータを書いてみる：

def file_line_reader_generator(file_path):
    """ファイルの行を返すジェネレータ"""
    with open(file_path, encoding="utf-8") as in_file:
        for line in in_file:
            yield line

イテレータオブジェクトを取得する：

file_lines = file_line_reader_generator(in_file_path)

イテレートしてみる：

for line in file_lines:
    print(line)

ちゃんと1行ずつ読めてる。

ここまで学んだことを踏まえて、

指定行数でファイルを分割して出力するスクリプトを大胆に書きなおしてみた：

# -*- coding: utf-8 -*-

import os.path;

def file_line_reader_generator(file_path):
    """ファイルの行を返すジェネレータ"""
    with open(file_path, encoding="utf-8") as in_file:
        for line in in_file:
            yield line

def output_file_from_generator(output_file_path, file_line_generator, max_lines_per_file):
    """ 
    行ジェネレータから行を読み込んで行数がmax_lines_per_filesに達するまで
    output_file_pathファイルに出力する
    """
    is_line_left_break = False
    with open(output_file_path, mode="w", encoding="utf-8") as out_file:
        for i,line in enumerate(file_line_generator):
            out_file.write(line)
            if i + 1 >= max_lines_per_file:
                is_line_left_break = True
                break

    return is_line_left_break

def split_file(in_file_path, output_dir_path, max_lines_per_file):
    """ 
    in_file_pathのテキストファイルをmax_lines_per_file行ごとに分割して
    output_dir_pathに出力する
    """
    file_name = os.path.basename(in_file_path)
    in_file_name_parts = file_name.split(".")
    out_file_name_template = "{0}_%d.{1}".format(*in_file_name_parts)

    file_lines = file_line_reader_generator(in_file_path)

    is_line_left = True
    file_index = 1 
    while (is_line_left):
        output_file_path = os.path.join(output_dir_path, out_file_name_template % file_index)
        is_line_left = output_file_from_generator(output_file_path, file_lines, max_lines_per_file)
        file_index = file_index + 1 

in_file_path = "＜分割前のファイルのパス＞"
max_lines_per_file = 1000
output_dir_path = "＜分割したファイルを出力するディレクトリのパス＞"
split_file(in_file_path, output_dir_path, max_lines_per_file)