Pythonでネストが深くなって読みにくい

はじめに

Pythonをかいていて、「ネストが深くて読みにくいなぁ」ってよく思うんですよね。
具体的には以下のような場合です。

for group in groups:
    for member in group.members:
        for article in member.articles:
            for tag in article.tags:
                """
                処理いろいろ
                """

もしこれが以下のようになったらステキじゃないですか？

for group, member, article, tag in func(groups):
    """
    処理いろいろ
    """

でも「そもそもfor文ってどうやったら自分で制御できるの？」って疑問もあると思うのでそこから解説していきたいとおもいます。

実装

Pythonを含め一部のプログラミング言語にはイテレータ(iterator)とジェネレータ(generator)が存在します。
Pythonのものについては以下の記事が参考になります。
Pythonのイテレータとジェネレータ

今回はこのジェネレータを用いて実装していきます。

まず軽く動作を確認します。

def my_generator():
    yield 10
    yield 20
    yield 30

for num in my_generator():
   print(num)

"""出力結果
10
20
30
"""

yieldで渡した値が次々にfor文の引数になっていることがわかります。

これを踏まえてサンプルのgroupsを渡したら、tagまで展開してforの引数に渡してくれる関数を作ります。

def group_generator(groups):
    for group in groups:
        for member in group.members:
            for article in member.articles:
                for tag in article.tags:
                    yield group, member, article, tag


for group, member, article, tag in group_generator(groups):
    print(group)
    print(member)
    print(article)
    print(tag)

これでネストの浅いソースコードになりました。
では更にこれを一般化していきます。
group_generatorを一般化するにあたって、２つ目のfor文以降、ある値から次のIterableなオブジェクトを取り出し、for文で回すという処理が共通しており、再帰的になっていることがわかります。またyeildしている箇所を再帰関数のベースケースとすることでうまく実装できそうです。

ある値から次のIterableなオブジェクトを取り出す処理をiterable_creators、ベースとなるオブジェクトをbase_iterable、再帰的な部分をnest_recursive関数として実装すると以下のようになります。
またジェネレータ関数を関数内部で呼び出す場合は、yield fromを用います。

def nests(base_iterable, *iterable_creators):
    def nest_recursive(nested_items: Tuple, current_item, level: int):
        try:
            iterable_creator = iterable_creators[level]
            next_iterable = iterable_creator(current_item)
            for next_item in next_iterable:
                new_stuck = (*nested_items, next_item)
                yield from nest_recursive(new_stuck, next_item, level + 1)
        except IndexError:
            yield nested_items

    for root_item in base_iterable:
        yield from nest_recursive((root_item,), root_item, level=0)

では実際に使ってみましょう。

iterable_creators = (lambda group: group.members, lambda member: member.articles, lambda article: article.tags)
for group, member, article, tag in nests(groups, *iterable_creators):
    print(group)
    print(member)
    print(article)
    print(tag)

ネストがとれて読みやすくなりましたね。
実用性があるかどうかは怪しいですが、ジェネレータの良い勉強になりました。