Pythonでネストしたfor文の途中で抜ける（ジェネレータの出番です）

はじめに

取り繕ってばかりの自分が 誰よりも嫌いなんだ♪
『未完成のポラリス』よすぎる、nikkieです。

みんなのPython勉強会#108のご参加、ありがとうございました。
懇親会で出た話題の1つに、アンサーブログを書きます。

目次

• はじめに
• 目次
• ネストしたfor文の途中で抜けたい
• 「Breaking out of nested loops with generators」より
• 気づき：ジェネレータはデータの流れ
• 終わりに

ネストしたfor文の途中で抜けたい

やりたいことは、例えば [[1, 2], [3, 4], [5, 6]] というリストから 1, 2, 3, と取り出していくのですが、3で全てのfor文を抜けたい（4以降は処理しない）というものです。

for sublist in [[1, 2], [3, 4], [5, 6]]:
    for element in sublist:
        print(element)
        # TODO element == 3 のときに2重のfor文を抜けたい

ジェネレータの出番という記事を思い出しました。

Breaking out of nested loops with generatorshttps://t.co/XxyCa9WJXY
[[1, 2], [3, 4], [5, 6]] から
1, 2, 3, と取り出すが、3でbreakしたい
2重のループだと書くのが難しいが、要素を1つずつ返すジェネレータ(*)を用意し、それを回すループで3ならbreakすればよいとのこと
(*)itertoolsにもある

— nikkie / にっきー 技書博 け-04 Python型ヒント本 (@ftnext) 2024年7月17日

「Breaking out of nested loops with generators」より

2重のfor文をジェネレータ関数¹とします。
これだけです。

def elements_from_sublists(nested_list):
    for sublist in nested_list:
        for element in sublist:
            yield element

for element in elements_from_sublists([[1, 2], [3, 4], [5, 6]]):
    print(element)
    if element == 3:
        break

上記のコードを1つのスクリプトにまとめて実行

% python -V
Python 3.12.6
% python script.py
1
2
3

なお、ジェネレータ関数を自分で書かずにitertools.chain.from_iterable()も紹介されています²。
https://docs.python.org/ja/3/library/itertools.html#itertools.chain.from_iterable

from itertools import chain

for element in chain.from_iterable([[1, 2], [3, 4], [5, 6]]):
    print(element)
    if element == 3:
        break

気づき：ジェネレータはデータの流れ

（ここの見出しの「ジェネレータ」はすぐ後に登場する「ジェネレータイテレータ」を指して使っています）

Python用語集より、「イテレータ」とは
https://docs.python.org/ja/3/glossary.html#term-iterator

データの流れを表現するオブジェクトです。

組み込み関数next()に渡すと「流れの中の要素を一つずつ返」す、とあります。
この流れという見方はなるほどと思います。

別の記事で書いているのですが、ジェネレータ関数はジェネレータイテレータを返します。

ジェネレータイテレータはイテレータ³の1種です（is-a関係）。
つまり、データの流れなんですよ！

2重のfor文を切り出したジェネレータ関数、これはデータの流れ"だけ"を表していると気づきました！

def elements_from_sublists(nested_list):
    for sublist in nested_list:
        for element in sublist:
            yield element

「ネストしたfor文の途中で抜けたい」は、データの流れを作るコード（for文）と抜けるためのロジックを一緒に考えているとなかなか実現できない⁴のですが、データの流れ（ジェネレータ）を切り出すだけで不思議なことに簡単に実現できるんです！

終わりに

ネストしたfor文の途中で抜けたいときには

• ネストしたfor文の部分をまずジェネレータ関数にしよう（紹介記事参照）
• ジェネレータ関数の返り値をfor文で反復して、抜けたい条件を書けば実現できる！

この記事を書いてみて、「ジェネレータ関数の返り値は、データの流れ（イテレータ）」という点は、見方が少し深まったように思います。
1つ1つの要素を取り出すロジック（今回はネストしていた）と、反復を制御するロジックを切り分けたんだ！

mathspp.com さん、示唆のある記事をありがとうございます！