Pythonのyieldわからないを解消してみる

はじめに

Pythonには、関数の中で値を一度に全て返すのではなく、1つずつ「遅延して」返すためのキーワードとしてyieldがあります。
yieldを使うことで、ジェネレータという特殊な種類の反復可能なオブジェクトを作成できます。

本記事では、以下を解説します

1. ジェネレーターの基本的な仕組み
2. メモリ効率の良さをCSV読み込みで比較検証

yieldとは？

通常の関数では、returnを使って値を返します。一方、yieldを使うと関数は値を「生成」し、その時点での状態を保存し、このような関数はジェネレータ関数と呼びます。
ジェネレータ関数は以下の特徴を持ちます

1. 一度に全ての値を計算しない（メモリ効率が良い）
2. 値を1つずつ取り出せる
3. 実行途中の状態を覚えていて、次回呼び出すと中断した場所から再開します

基本的な例

yieldを使用しない場合とyieldを使った実装を比較してみます。

yieldを使用しない場合

def get_numbers():
    return [1, 2, 3, 4, 5]

numbers = get_numbers()
for num in numbers:
    print(num)

yieldを使ったジェネレータ関数

def generate_numbers():
    for i in range(1, 6):
        yield i

numbers = generate_numbers()  # ジェネレータオブジェクトを作成
for num in numbers:  # yieldを呼び出し
    print(num)

実行の流れ

yieldは関数を一時停止し、次に呼び出されたときに中断した場所から再開します。

1. 最初に関数を呼び出しても、実行は開始されません（ジェネレータオブジェクトを作成するだけ）
2. forループやnext()で値を取得するたびにyieldが呼び出されます
3. 最後の値を返した後、ジェネレータはStopIteration例外を発生させて終了します

def simple_generator():
    print("開始")
    yield 1
    print("中間地点")
    yield 2
    print("終了")
    
gen = simple_generator()
print(next(gen))  # 開始, 1
print(next(gen))  # 中間地点, 2
print(next(gen))  # 終了, StopIteration例外

実際にメモリ効率が良いのか調べてみた

ジェネレータのメモリ効率を、5万行のデータが記載されたCSVファイルを処理する例で比較します。

ジェネレーターを使用した場合のサンプルコード

import csv
import os
import psutil

def memory_usage():
    """現在のプロセスのメモリ使用量をMB単位で返す"""
    process = psutil.Process(os.getpid())
    mem = process.memory_info().rss / 1024 / 1024  # MBに変換
    return mem

def read_large_csv(file_path):
    """CSVファイルを1行ずつ読み込むジェネレーター"""
    with open(file_path, mode='r', encoding='utf-8') as f:
        reader = csv.reader(f)
        for row in reader:
            yield row  # 1行ずつ返す

# メモリ使用量を測定
print("ジェネレーターを生成前のメモリ使用量: {:.2f} MB".format(memory_usage()))
large_generator = read_large_csv("サンプル.csv")
print(next(large_generator))  # 1行目を取得
print(next(large_generator))  # 2行目を取得
print("ジェネレーターを生成後のメモリ使用量: {:.2f} MB".format(memory_usage()))

リストを使用した場合のサンプルコード

def read_large_csv(file_path):
    """CSVファイルをリストで読み込む"""
    with open(file_path, mode='r', encoding='utf-8') as f:
        reader = csv.reader(f)
        rows = list(reader)  # 全データをリストにロード
    return rows

# メモリ使用量を測定
print("リストを生成前のメモリ使用量: {:.2f} MB".format(memory_usage()))
large_list = read_large_csv("サンプル.csv")
print(large_list[0])  # 1行目を取得
print(large_list[1])  # 2行目を取得
print("リストを生成後のメモリ使用量: {:.2f} MB".format(memory_usage()))

結果

処理方法	生成前のメモリ使用量	生成後のメモリ使用量
ジェネレータ	約27 MB	約28 MB
リスト	約25 MB	約474 MB

ジェネレータすごい！

まとめ

今回はジェネレータについてまとめてみました。
大量データを扱う時には、ぜひyieldを活用したジェネレータを試してみてください！