Python の yield で値を返す

yield は、関数から値を返しつつ、その状態を保持する特別なキーワードです。return とは異なり、関数の実行を一時停止して、次に呼び出されたときに続きから再開できます。

yield と return の違い

return は関数を終了させて値を返しますが、yield は値を返した後も関数の状態を保持します。

return

関数を終了する。状態は破棄される。1回だけ値を返せる。

yield

関数を一時停止する。状態は保持される。何度でも値を返せる。

yield の基本

yield を含む関数は「ジェネレータ関数」となり、呼び出すとジェネレータオブジェクトを返します。

def my_generator():
    yield 1
    yield 2
    yield 3

gen = my_generator()
print(type(gen))  # <class 'generator'>

print(next(gen))  # 1
print(next(gen))  # 2
print(next(gen))  # 3

状態が保持される仕組み

yield で一時停止した関数は、ローカル変数の値をすべて保持しています。next() を呼ぶと、停止した場所から実行が再開されます。

def counter():
    print("開始")
    n = 0
    while True:
        n += 1
        print(f"yield前: n={n}")
        yield n
        print(f"yield後: n={n}")

gen = counter()
print(next(gen))
# 開始
# yield前: n=1
# 1

print(next(gen))
# yield後: n=1
# yield前: n=2
# 2

for 文で使う

ジェネレータは for 文で簡単に使えます。

def countdown(n):
    while n > 0:
        yield n
        n -= 1

for num in countdown(5):
    print(num)
# 5, 4, 3, 2, 1

yield を使うメリット

yield を使うと、大量のデータを一度にメモリに載せる必要がなくなります。

# リストを返す場合(全データがメモリに載る)
def get_squares_list(n):
    result = []
    for i in range(n):
        result.append(i ** 2)
    return result

# yield を使う場合(1つずつ生成)
def get_squares_gen(n):
    for i in range(n):
        yield i ** 2

100万件のデータを扱う場合、リストは全件をメモリに保持しますが、ジェネレータは1件ずつ処理するのでメモリを節約できます。