内包表記とyieldそしてジェネレータについて

内包表記とyieldそしてジェネレータについて

こんにちは、Yumihikiです。 Pythonを使う中で抑えておきたいテクニック、概念を学んだのでアウトプットします。

この記事で学べること

  • 内包表記という概念
  • yieldを使うメリット
  • ジェネレータとは何か

今回の題材

次のようなコードを題材にしてみます。

import csv


def csv_to_list(csv_file_path: str) -> list:
    """CSVファイルを読み込んで返却する

    :param csv_file_path: CSVファイルのパス
    :return: CSVファイルの中身
    """
    with open(csv_file_path, "r", encoding="utf-8") as f:
        rows = []
        for row in csv.reader(f):
            rows.append(row)
        return rows

このコードは、引数として渡されたCSVファイルを読み込み、その中身をリストとして返すものです。 良くある書き方だと思いますが、改善の余地があるので、改善例をいくつか見ていきましょう。

リストの内包表記を使ってみる

先程のコードは

  • CSVファイルを受け取る
  • リストにする

という単純な変換を行っているだけです。

そこで、次のように書き換えてみます。

import csv


def csv_to_list(csv_file_path: str) -> list:
    """CSVファイルを読み込んで返却する

    :param csv_file_path: CSVファイルのパス
    :return: CSVファイルの中身
    """
    with open(csv_file_path, "r", encoding="utf-8") as f:
        return [row for row in csv.reader(f)]

1行で書けてスッキリしました。 単純なリストへの展開であれば、内包表記を使った方が行数も減り、可読性も上がると思うので積極的に使っていくべきだと思います。

個人的に「最初にrows = []を宣言するのが面倒だな」「appendした結果を返却しているだけなので改善できないかな」と思っていたので、内包表記を積極的に利用するようにしています。

yieldを使ってジェネレータを返却する

さて、先程まではリストの内包表記を利用して、CSVファイルを読み込んでいました。 しかし、リストを使うことによるデメリットがあるのはご存知でしょうか。

そもそもの話として、変数に値を格納する際、メモリを使用しています。 データが膨大になればなるほどメモリを使用することになり、 パフォーマンスの悪いプログラムになってしまいます。

ではそれを改善した次のコードを見てみましょう。

import csv
from collections.abc import Generator


def csv_to_generator(csv_file_path: str) -> Generator:
    """CSVファイルを読み込んで返却する

    :param csv_file_path: CSVファイルのパス
    :return: CSVの行
    """
    with open(csv_file_path, "r", encoding="utf-8") as f:
        for row in csv.reader(f):
            yield row

これはyield式を使って書き換えたコードです。

返り値もlistからGeneratorに変わりました。 yieldを利用することで通常の関数がジェネレータ関数というものになり、Generatorを返却するためです。

yield 式及び文は generator を定義するときに、その本体内でのみ使うことが出来ます。関数定義内で yield を使用することで、その定義は通常の関数でなくジェネレータ関数になります。

https://docs.python.org/ja/3/reference/simple_stmts.html#the-yield-statement

前までのコードはリストを一括で返却するものでしたが、このコードは1行だけ返却することによって、メモリを節約できるようになります。

ジェネレータは公式リファレンスの用語集で次のように記載されています。

(ジェネレータ) generator iterator を返す関数です。 通常の関数に似ていますが、 yield 式を持つ点で異なります。 yield 式は、 for ループで使用できたり、next() 関数で値を 1 つずつ取り出したりできる、値の並びを生成するのに使用されます。

ジェネレータ

ジェネレータを利用する場合の注意点

ジェネレータは「1回」しか利用できません。

def gen():
    yield 1
    yield 2
    yield 3

g = gen()

# 1回目なので出力される
for i in g:
    print(i)

# 2回目なので何も出力されない
for i in g:
    print(i)

また、リストのようにインデックスでアクセスすることは出来ません。

def gen():
    yield 1
    yield 2
    yield 3

g = gen()

# TypeError: 'generator' object is not subscriptable が発生する
g[0]

補足: ジェネレータの内包表記もできます

今回のCSV読み込みではファイルが開きっぱなしになってしまうため利用できませんが（=withステートメントを使う必要がある）、ジェネレータの内包表記もできます。

list_ = [100, 200, 300, 400, 500]

# リストの内包表記を利用する場合
sum([price for price in list_])

# ジェネレータの内包表記を利用する場合
sum((price for price in list_), 0)

# １つだけなら省略することもできる
sum(price for price in list_)

ポイントはリストの内包表記で利用していた[]を利用するのではなく()で囲むことです。 これだけでジェネレータの内包表記になります。

また、代入する際に yield 式が１つだけの場合は括弧の省略が可能です。

yield 式が代入文の単独の右辺式であるとき、括弧は省略できます。

6.2.9. Yield 式より

最後に

今回はこのようなことを学びました。 これまではyieldを使ったらlistより便利なんでしょ？くらいの解像度の理解でしたが

  • リストより効率的にメモリを利用できる
  • 可読性が良くなる
  • リストと同じようにfor文を回せるが、1回しか利用できない、インデックスでアクセスできない

という理解まで深める事ができました。

特に、何気なくlistを返却することは出来るだけ控えようと思いました。 最初の頃に学んだ「変数はデータを格納する箱」のような認識しか持っていなかったので、きちんと裏側で何をしているか意識する必要があるなと痛感しました。

こうしたことを学ぶのは楽しいので、引き続き学んでいこうと思います。

参考