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

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