みーのぺーじ

みーが趣味でやっているPCやソフトウェアについて.Python, Javascript, Processing, Unityなど.

Pythonのジェネレーター内包表記にハマった

バグの原因が分かるまでに2時間もかかったので,自分への戒めの記事です.

話を簡単にするために,4つのフルーツの名前を扱うプログラムを例に用います.

fruits = set(["apple", "banana", "mango", "orange"])
print(sorted(fruits))

これをPython3.8で実行すると,

['apple', 'banana', 'mango', 'orange']

となります.

大文字と小文字の表記にする関数

受け取った引数に大文字と小文字の表記を追加する'upper_and_lower(items)'関数を作成します.

def upper_and_lower(items):
    r = set(items)
    r.update(u.upper() for u in items)
    r.update(u.lower() for u in items)
    return r

upper_and_lower_fruits = upper_and_lower(fruits)
print(sorted(upper_and_lower_fruits))

これを実行すると,

['APPLE', 'BANANA', 'MANGO', 'ORANGE', 'apple', 'banana', 'mango', 'orange']

となります.うまく動作しています.

Pythonのsetのupdate()関数は,listはもとより,setやgeneratorでも引数に指定できるので,もしもitemsが巨大なリストでも問題がないように,関数内部でジェネレーターを利用することで,メモリー消費を抑えるよう工夫しています.

大文字はじめの表記も追加する

この関数を使って,大文字はじめの表記("Apple"など)を追加してみます.引数を少し変えます.

upper_lower_and_capitalize_fruits = upper_and_lower(
    u.capitalize() for u in fruits
)
print(sorted(upper_lower_and_capitalize_fruits))

期待していた出力は,

['APPLE', 'Apple', 'BANANA', 'Banana', 'MANGO', 'Mango', 'ORANGE', 'Orange', 'apple', 'banana', 'mango', 'orange']

でしたが,実際は

['Apple', 'Banana', 'Mango', 'Orange']

となります.

当たり前の話なのですが,ジェネレーターは一度しか値を取得できません.今回の'upper_and_lower(items)'関数はitems変数を3回参照しているので,2回目以降の参照で正しく処理できない問題が生じています.

修正案1

'upper_and_lower(items)'関数の引数をジェネレーターではなくsetとすることでこの問題を解消することができます.しかし,引数に制約が出るのがデメリットです.

upper_lower_and_capitalize_fruits = upper_and_lower(
    {u.capitalize() for u in fruits}
)
print(sorted(upper_lower_and_capitalize_fruits))

修正案2

'upper_and_lower(items)'関数内で引数を参照する回数を 1回にすることでこの問題を解消することができます.修正案1とは異なり,引数にはジェネレーターを利用することができるのがメリットです.

def upper_and_lower(items):
    r = set()
    for u in items:
        r.add(u)
        r.add(u.upper())
        r.add(u.lower())
    return r

ジェネレーターは便利なので,ついリストと同じような扱いをしてしまいますが,自分で作った関数の引数がジェネレーターに対応しているかを意識するべきでした.