pythonの generator について、yieldはもともとたまに使ってたのですが

  • sendやreturnもできること
  • 非同期処理でも使えること

を知ったので、軽く記事を書いておこうと思います。

ジェネレータ（Generator）

ジェネレータ関数は、値を返す代わりにジェネレータ（一連の値を返す特殊なイテレータ）を返す関数で、yieldキーワードを含む関数はジェネレータ関数です。
ジェネレータ関数の中でyieldが呼び出されると、一旦ジェネレータ関数は一時停止して呼び出しの処理が実行されます。次にジェネレータの__next__()が呼び出される（for文やnext(gen)など）と、ジェネレータ関数の処理が再開されます。

Listを返して処理できるケースであれば、ジェネレータ関数にする必要はないと思いますが、

  • 値をyield（生成）した時の状態で何か処理をしたい場合（Listにすると値を生成した際の状態は失われている）
  • Listの件数が多すぎてリソース（メモリ）が足りなくなるので、逐次処理したい場合

など、ジェネレータ関数を利用した方が良いケースがあります。

yield

一番シンプルなyieldだけを行うサンプルです。

import time
from typing import Generator

def main():
    for i in generate():
        print(f'get. i: {i}')
        time.sleep(1)

def generate() -> Generator[int, None, None]:
    print('generate started.')
    for i in range(3):
        print(f'before yield. i: {i}')
        yield i
        print(f'after yield. i: {i}')
    print('generate finished.')

if __name__ == "__main__":
    main()

実行すると以下のログが出力されました。

generate started.
before yield. i: 0
get. i: 0
after yield. i: 0
before yield. i: 1
get. i: 1
after yield. i: 1
before yield. i: 2
get. i: 2
after yield. i: 2
generate finished.

ポイントは以下です。

  • before yield => get => after yield の順でログが出力されていることから、yieldした時点でジェネレータ関数は一時停止し、呼び出し元に処理がうつっている
    • 呼び出し元ではgetの後1秒sleepしているがその間もジェネレータ関数は停止したまま
    • 呼び出し元で次のloopに入った時（内部的にはgenerator.next()が呼び出されたタイミング）で、ジェネレータ関数側の処理が再開される
  • generate startedとgenerate finishedはそれぞれ1回しか出力されてないことから、generate()は一度しか呼び出されない

send

最近知ったのですが、yieldして一時停止したジェネレータ関数に、呼び出し元からgenerator.send()を使って値を渡すことができます。
ジェネレータ関数から受け取った後呼び出し元で処理した結果で、ジェネレータ関数の挙動を変えたい時に使うものだと思います。

import time
from typing import Generator, Optional
from random import randint

def main():
    gen: Generator = generate()  # Generatorオブジェクト生成
    i: int = 0
    while True:
        try:
            print(f'before next. i: {i}')
            i = next(gen)
            print(f'after next. i: {i}')
            time.sleep(1)
            send_value: int = randint(1, 10)  # 1〜10のランダムな整数を取得
            print(f'before send. i: {i}, send_value: {send_value}')
            gen.send(send_value)
            print(f'after send. i: {i}, send_value: {send_value}')

        except StopIteration:
            print('StopIteration raised.')
            break

def generate() -> Generator[int, int, None]:
    print('generate started')
    i: int = 1
    while True:
        print(f'before yield. i: {i}')
        received_value: Optional[str] = (yield i)  # next()でもsend()でもここに来る。next()だとreceived_valueはNone
        print(f'after yield. i: {i}, received_value: {received_value}')
        if received_value is not None:
            i = i + received_value  # Noneじゃなければiに受け取った値を加算
        if i > 10:  # 10を超えたらloop終了
            break
    print('generate finished')

if __name__ == "__main__":
    main()

実行すると以下のログが出力されました。

before next. i: 0
generate started
before yield. i: 1
after next. i: 1
before send. i: 1, send_value: 6
after yield. i: 1, received_value: 6
before yield. i: 7
after send. i: 1, send_value: 6
before next. i: 1
after yield. i: 7, received_value: None
before yield. i: 7
after next. i: 7
before send. i: 7, send_value: 8
after yield. i: 7, received_value: 8
generate finished
StopIteration raised.

ポイントは以下です。

  • for文ではsend()は使えないので、Generatorオブジェクトを取得する必要がある
  • 最初のbefore nextの後にgererate startedが出力されているため、Generatorオブジェクトを取得しても、ジェネレータ関数内部の処理は動いてない
  • ジェネレータ関数はyieldで停止した後、next(generator)でもgenerator.send()でも処理が再開される
  • send()の場合は値を受け取ることができるが、next()の場合はNoneを受け取る
  • ジェネレータ関数の処理が完了したら、StopIterationがraiseされる

return

ジェネレータ関数も、returnすることで関数自体の戻り値を返すことができます。ただし、ジェネレータ関数はもともとGeneratorオブジェクトを返却するためreturn valueを実行すると、StopIteration(value)をraiseするという少し特殊な動きになります。

import time
from typing import Generator

def main():
    gen: Generator[int, None, str] = generate()
    i: int = 0
    while True:
        try:
            print(f'before next. i: {i}')
            i: int = next(gen)
            print(f'after next. i: {i}')
            time.sleep(1)
        except StopIteration as e:
            print(f'StopIteration raised. return_value: {e.value}')
            break

def generate() -> Generator[int, None, str]:
    print('generate started.')
    for i in range(3):
        yield i
    print('generate finished.')
    return '##return_value##'

if __name__ == "__main__":
    main()

実行すると以下のログが出力されました。

before next. i: 0
generate started.
after next. i: 0
before next. i: 0
after next. i: 1
before next. i: 1
after next. i: 2
before next. i: 2
generate finished.
StopIteration raised. return_value: ##return_value##

ポイントは以下です。

  • ジェネレータ関数でreturnした値は、exceptでキャッチして受け取る必要がある

非同期ジェネレータ（AsyncGenerator）

スクレイピング処理などジェネレータ関数内で非同期処理を使いたいケースもあると思います。調べたところAsyncGeneratorというやつがいて、通常のGeneratorとほぼ同じ使い方で非同期ジェネレータ関数を書くことができました。
なお、通常のジェネレータと違って非同期ジェネレータはジェネレータ関数自体の戻り値を返すことはできません。

yield

単純にyieldするだけのサンプルです。以下を確認しています。

import asyncio
from typing import AsyncGenerator

async def main():
    print('====== async forの例 ======')
    async for i in async_generate():
        print(f'started loop. i: {i}')
        await asyncio.sleep(1)

    print('====== generatorを取得してwhileで回す例 ======')
    gen: AsyncGenerator[int, None] = async_generate()
    i: int = 0
    while True:
        try:
            print(f'before anext. i: {i}')
            i = await gen.__anext__()
            print(f'after anext. i: {i}')
        except StopAsyncIteration:
            break

async def async_generate() -> AsyncGenerator[int, None]:
    print('async_generate started.')
    for i in range(3):
        print(f'before yield. i: {i}')
        yield i
        print(f'after yield. i: {i}')
        await asyncio.sleep(1)
    print('async_generate finished.')

if __name__ == "__main__":
    asyncio.get_event_loop().run_until_complete(
        main()
    )

実行すると以下のログが出力されました。

====== async forの例 ======
async_generate started.
before yield. i: 0
started loop. i: 0
after yield. i: 0
before yield. i: 1
started loop. i: 1
after yield. i: 1
before yield. i: 2
started loop. i: 2
after yield. i: 2
async_generate finished.
====== generatorを取得してwhileで回す例 ======
before anext. i: 0
async_generate started.
before yield. i: 0
after anext. i: 0
before anext. i: 0
after yield. i: 0
before yield. i: 1
after anext. i: 1
before anext. i: 1
after yield. i: 1
before yield. i: 2
after anext. i: 2
before anext. i: 2
after yield. i: 2
async_generate finished.

ポイントは以下です。

  • async defで関数定義するだけで、非同期ジェネレータ関数を作成できる
  • async forで非同期ジェネレータをイテレートすることができる
  • もちろん、非同期ジェネレータオブジェクトを受け取って、await AsyncGenerator.__anext__()でイテレートすることもできる
    • anext(AsyncGenerator)やnext(AsyncGenerator)は存在しない

asend

非同期ジェネレータでも呼び出し元から、ジェネレータ関数に値を通知することができます。

import asyncio
from typing import AsyncGenerator
from random import randint

async def main():
    gen: AsyncGenerator[int, None] = async_generate()
    i: int = 0
    while True:
        try:
            print(f'before anext. i: {i}')
            i = await gen.__anext__()
            print(f'after anext. i: {i}')
            send_value: int = randint(1, 10)
            print(f'before anext. i: {i}, send_value: {send_value}')
            await gen.asend(send_value)
            print(f'after anext. i: {i}, send_value: {send_value}')
        except StopAsyncIteration:
            break

async def async_generate() -> AsyncGenerator[int, None]:
    print('async_generate started.')
    for i in range(10):
        print(f'before yield. i: {i}')
        received_value: int = (yield i)
        print(f'after yield. i: {i}, received_value: {received_value}')
        if received_value is not None and received_value > 5:
            print('finish async_generate because received_value > 0.')
            break  # 受け取った値が5より大きかったら終了する
        await asyncio.sleep(1)
    print('async_generate finished.')

if __name__ == "__main__":
    asyncio.get_event_loop().run_until_complete(
        main()
    )

実行すると以下のログが出力されました。

before anext. i: 0
async_generate started.
before yield. i: 0
after anext. i: 0
before anext. i: 0, send_value: 7
after yield. i: 0, received_value: 7
finish async_generate because received_value > 0.
async_generate finished.

ポイントは以下です。

  • 非同期ジェネレータではawait AsyncGenerator.asend()で通知する（awaitを忘れずに）
  • 非同期ジェネレータ関数と呼び出し元間での処理の流れは、通常のジェネレータのsend()を使うケースと同じ

さいごに

ジェネレータは、yieldした時に一旦処理を停止して呼び出し元の処理が完了したら次の処理を行えるところがとても良くて、ジェネレータを覚えてから、for文が深くなっていやだなーと思っていた箇所を、ジェネレータ関数にして階層を少なくして可読性がだいぶ上がったように感じます（多少読み手のリテラシーが求められる部分はありますが）。

とくにヘッドレスブラウザを使ったスクレイピングみたいな、リストとして何かを返すことができないような処理とは相性がよかったです。

かなり便利な機能なので今後も使い倒していこうと思います。

関連記事
  1. 標準モジュールunittestでpythonのテストを書く
  2. Pythonでasyncioを使った非同期処理
  3. Pythonにおけるマルチプロセスでの状態管理
  4. Scrapyの使い方
  5. pythonのTips