Python の defaultdict には注意しよう

こんにちは。

Python には標準的なデータ型に加えていくつかの便利な型が用意されています。特に collections モジュールには、以前記事に書いた deque や、今回取り上げる defaultdict が格納されています。

通常、辞書型のデータ d にデフォルト値を設定しながら値を詰めていく場合、d.setdefault(key, default_value) といったやり方をするのではないかと思います。このようにしない場合は、例えば d に保存されていないキー k で辞書にアクセスしてしまうと、KeyError が出てしまうのですね。

同様の用途で、より柔軟かつ（おそらく）高速な処理が実現できるのが defaultdict です。

使い方としては、例えば

from collections import defaultdict

dd = defaultdict(lambda: dict())

dd['k1']  # dict(). 存在しないキーを入れるとデフォルト値が返る
dd['k2']['k2_k1'] = 'v1'    # デフォルト値の dict() に値を代入できる
dd['k2']  # {'k2_k1': 'v1'}

というようになります。defaultdict の引数になっている無名関数は引数をとりません。無名関数の返り値を defaultdict にすれば複数層の辞書を返すこともできるし、リストを返してもよいです。setdefault を使う場合とは異なり、値を代入するたびにデフォルト値を指定する必要もありません。このように、defaultdict は、データを保管する際に便利に使うことができます。

一方で、データを取り出す際には注意が必要、ということに最近気がつきました。上に書いたように defaultdict は存在しないキーに対してもデフォルト値を返してしまうので、たとえコードの中で意図しないキーを入れた場合であっても、KeyError を出さないがゆえに隠れたバグの原因になることがあります。

私が実際にやってしまったことがある例としては、こんなことがありました。ある defaultdict dd が「キー：辞書」のペアを格納しているのですが、このキーのリストは、本来別の dict d (通常の辞書) と共通しているようにすべき状況でした。ところが、意図しないキーを用いて

len(dd["hoge"])  # 0

と dd にアクセスしてしまったがために、dd の内部では defaultdict( at 0x1057e0ea0>, {'hoge': {}}) となり、意図しないキーが保存されてしまいました。この状態で for k in dd: としてキーに対してループを回すと、特に辞書 d においては該当するキーがが存在しないため KeyError が出ることとなってしまいます。

これを防ぐためには、必要なデータを保管し終えた時点で

dd = dict(dd)

として通常の辞書にしてしまうといいでしょう。あるデータ準備用の関数内で用意した defaultdict を return するのなら、return 時に変換しておくと安心です。こうすれば、従来通り存在していないキーに対してエラーが出るようになります。

これには良い副産物もあって、辞書化したデータは pickle で永続化することができるようになります。というのは、lambda 関数が pickle 化できないためですね。試しに defaultdict を pickling しようとすると、

import pickle
from collections import defaultdict

dd = defaultdict(lambda: dict())
dd['k1']['k1_k1'] = 'v1'
with open('dd.pkl', 'wb') as f:
    pickle.dump(dd, f)
# PicklingError: Can't pickle <function <lambda> at 0x1106096e0>: it's not found as __main__.<lambda> 

となってしまうはずです。

最近、pandas を使おうにも使えない微妙なデータがあり、仕方なく defaultdict を使ったところ、どうもエラーは出ないけど結果は明らかにおかしい、どうしたものか、あ変な空辞書が入っている、みたいな状況になってしまって、教訓として得られたのが今回の記事になっています。みなさんもお気をつけて。