2013年9月18日

PythonのOrderedDictで循環参照を扱えるようにする

はじめに

Python の OrderedDict は挿入の順番を保持してくれる順序つきの辞書で、ある場面ではとっても有用なのですが、組み込みのdictと違って循環参照をうまく扱えません。たとえば、以下のように OrderedDictオブジェクトを作成し、要素に自分自身を代入してみます。 ここまでは良いのですが、ここで作成したOrderedDictオブジェクトを表示しようとすると無限再帰呼び出しに突入し、実行時エラーとなってしまいます。

何が問題か?

ordereddictのソースコードを見てみると、原因はすぐにわかります。 self.items()の呼び出しで、子要素を取り出して(%r)でフォーマット出力 していますが、この部分で再び__repr__が呼ばれるので、OrderedDictに循環参照が含まれていると無限の再帰呼び出しに突入してしまうのです。
循環参照を正しく扱うには、__repr__が呼び出された時のselfをどこかに記憶(メモ)しておいて、再び呼び出されたときに、既に呼び出されているかどうかをチェックする処理を追加する必要があります。 たとえば、以下のように実装するとよいでしょう。 この例では、スレッドローカル変数に、訪問済みかどうかをチェックするための辞書を作成しておき、すでに訪問済みの場合は「...」を返すようにしておいます。 実行してみると次のようになります。

pickle化に対応する

さて、循環参照データの表示の問題は解決したのですが、さらに、これをPickle化しようとすると、同じように実行時エラーとなります。 Pythonのpickle化処理では、オブジェクトに__reduce__()が実装されている場合、__reduce__()が呼び出されます。そこで、OrderedDictの__reduce__()の実装を見てみることにします。 局所変数のitemsに、OrderedDictが保持しているkey-valueの組みのリストを代入して、それをOrderedDictクラスとともに返しています。Unpickle時には、itemsを引数としてここで返されたクラスが呼び出されます。さらに原因を調べるために、Pythonのpickle化処理のソースコード(/usr/lib64/python2.6/pickle.py)を追ってみると、__reduce__()で返されたitemsに対して再帰的にpickle化処理が呼ばれるようになっています。ということで、循環参照を持つデータをpickle化しようとすると、無限の再帰呼び出しに突入することになります。
再び、pickle化処理のソースコードを追ってみると、メモを用いてオブジェクトの循環データ構造をきちんと扱うための仕組みがもともと実装されいていることがわかります。今回はこれをうまく活用して、この問題を解決することにします。先ほどのOrderedDictの__reduce__()の実装では、自分自身を要素として含む可能性のあるitemsをタプルの第2要素として返していました。pickle化のソースコードを調べてみると、ここで返した要素は、pickle化対象のオブジェクトをメモに記録する前に再帰的にpickle化していることがわかります。そこで、itemsをタプルの第2要素として返すのはあきらめ第3要素として返すことにします。__reduce__()が返す第3要素は、unpickle化のさいにオブジェクトの__setstate__()メソッドの引数として渡されます。幸いなことに、この第3要素は、メモに記録した後にpickle化されるため、自分自身をその要素に含んでいたとしても、メモに記録が残っているため再帰呼び出しが抑止されます。ということで、以下のように実装すれば問題は解決します。 実際にpickle/unpickle化をためしてみます。 ということで、循環参照データでもうまくpickle/unpickle化できるようになりました。

本当に問題は解決されたのか?

これで、めでたしめでたし、と終わりにしたいところなのですが、実はまだ問題が残っています。以下のように循環参照のあるデータ同士を比較すると、やっぱり無限に再帰呼び出しに突入してしまいます。 実はこの問題はOrderedDictだけのものではなく、Pythonの組み込みのdictでも同じ現象が発生します。循環参照を含むデータの構造比較を正しく扱うのは難しそうなので、今回はこれまでということにいたします。