Pythonのyield文を使って処理の進捗状況を返す

はじめに

Pythonの関数内でfor文を使って重たい処理を回すときに、yield文を使って進捗状況を呼び出し元に返す方法を解説します。

想定する状況：

ディープラーニングやデータ送受信などの時間が掛かる処理をfor文で実行する。
処理の進捗をTkinter, PyQt, Django, FlaskなどでGUI表示（プログレスバー等）したい。

進捗を表示する処理をfor文の中に入れることも実装の一つとして可能です。しかし、大規模なアプリケーションではGUIとビジネスロジックを分離することが一般的であるため、for文の中に直接GUIに関するコンポーネントを入れたくありません。

このGUIとビジネスロジックを分離する方法について、具体的な実施例が見つけられなかったため備忘録として残します。

検証環境

Python 3.11.6
tqdm 4.66.2

yield文で反復回数を返す

処理の進捗状況を返すには、Pythonのyield文を用います。例を以下に示します。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


from typing import Iterator

def my_func(n_iteration:int) -> Iterator[int]:
    for i in range(n_iteration):
        # ここに反復処理する重たい処理を書く
        yield i

if __name__=="__main__":
    for j in my_func(5):
        print(j)

実行結果は以下になります。my_func内のyieldで返した値が、変数jに格納されることが分かります。そのため、呼び出し元で、何回目のfor文ループが実行されているのか把握できます。

ここでyieldについて解説します。yieldはreturnと似ている構文で、関数の呼び出し元に値を返します。ただし、returnと異なる点として、yield文では「値を返した後、元の関数の処理に戻る」挙動となります。

tqdmを使用した例

進捗を表示できるライブラリtqdmを使用した例を以下に示します。wait_1sec関数では、forループを1回実行する度に1秒待ちます。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14


from typing import Iterator
import time
from tqdm import tqdm

def wait_1sec(n_iteration:int) -> Iterator[int]:
    for i in range(n_iteration):
        time.sleep(1)
        yield i

if __name__=="__main__":
    n_iter = 5

    for _ in tqdm(wait_1sec(n_iter), total=n_iter):
        pass

実行結果は以下になります（実際はプログレスバーが増える様子がアニメーションで表示されます）。

100%|██████████████████████████| 5/5 [00:05<00:00,  1.00s/it]

tqdmの仕様上、呼び出し元のforループが実行されるだけでプログレスバーが増えていくため、yield文でループ回数を返す意味はありません。ただし、GUIライブラリによってはループ回数が必要になるかもしれないため、ループ回数を返した方がGUIの変更に対してロバスト性があると思います。