Python で MultiIndex

土曜日に boost 勉強会の UStream 見てたら Python でも MultiIndex を使いたくなったので作ってみた．

ちなみに，Python3 です．

こんな風に使いたい!!

MultiIndex は何を index とするかをどんな風に表現するのかが肝なのかなぁって思う．

boost の場合は，型を指定したり，メンバーのアドレス指定したりといろいろめんどくさそう．

せっかく Python なのでそんな制限したくないし，どうせつっこむデータの型なんて同じなんだし，とういことで関数オブジェクトとかで指定するようにしてみた．

実際には↓な感じで使えるように作る．

import collections

class Employee:
    def __init__(self, name, age, sex):
        self.name = name;
        self.age = age
        self.sex = sex
        
    def __str__(self):
        return '({})'.format(', '.join([self.name, str(self.age), self.sex]))

    def __repr__(self):
        return str(self)

def multi_index_test():
    mdict = MultiIndex(name=Unique(lambda x: x.name, collections.OrderedDict),
                       age=NonUnique(lambda x: x.age),
                       sex=NonUnique(lambda x: x.sex))
    mdict.append(Employee('Mike', 20, 'male'))
    mdict.append(Employee('Jane', 21, 'female'))
    mdict.append(Employee('Bob', 21, 'male'))
    
    print(mdict['name']['Mike'])
    print(mdict['age'][21])

実装

Python は基本的に shallow copy なので得に気にすることなく，標準のコンテナにどんどんつっこむようにする．

問題は追加，削除時のいっせい更新だけど，それは Observer パターンで解決する．

そんなんでできたのが↓

class Index(object):
    def __init__(self):
        self._index_holder = None

    def _set_index_holder(self, index_holder):
        self._index_holder = index_holder

    def append(self, obj):
        self._index_holder._complete_add(obj)

    def pop(self, obj):
        self._index_holder._complete_pop(index)


class Unique(Index):
    def __init__(self, rule, container=dict):
        super(Index, self).__init__()
        self._make_key = rule
        self._contents = container()

    def __getitem__(self, k):
        return self._contents[k]

    def __contains__(self, obj):
        return self._make_key(obj) in self._contents

    def __iter__(self):
        return self._contents.values()

    def _add(self, obj):
        key = self._make_key(obj)
        self._contents[key] = obj

    def _erase(self, obj):
        key = self._make_key(obj)
        return self.contents.pop(key)


class NonUnique(Index):
    def __init__(self, rule, container=dict, drawer=list):
        super(Index, self).__init__()
        self._make_key = rule
        self._drawer = drawer
        self._contents = container()

    def __getitem__(self, k):
        return self._contents[k]

    def __contains__(self, obj):
        return obj in self[obj]

    def __iter__(self):
        for drawer in self._contents.values():
            for obj in drawer:
                yield obj

    def _add(self, obj):
        key = self._make_key(obj)
        self._contents.setdefault(key, self._drawer()).append(obj)

    def _erase(self, obj):
        key = self._make_key(obj)
        return self._contents[key].pop(obj)


class MultiIndex(object):
    def __init__(self, *args, **kwargs):
        self._catalog = {}
        for i, index in enumerate(args):
            index._set_index_holder(self)
            self._catalog[i] = index
        for i, index in kwargs.items():
            index._set_index_holder(self)
            self._catalog[i] = index

    def __getitem__(self, k):
        return self._catalog[k]

    def _complete_add(self, obj):
        for index in self._catalog.values():
            index._add(obj)

    def _complete_erase(self, obj):
        for index in self._catalog.values():
            index._pop(obj)

    def append(self, obj):
        self._complete_add(obj)

    def pop(self, obj):
        self._complete_pop(obj)

Python の場合，dict が hash map なので全部 dict を使っちゃう．

入れた順番を覚えておきたい場合も collections.OrderedDict っていうのがあるので，dict と collections.OrderedDict を使い分けれるように引数で実際に使用するコンテナを渡せるようにした．

NonUnique の場合，ある index で検索した結果がただのリストでいい場合と，さらに MultiIndex の方がいい場合とがあると思うので，そこも引数で選択できるようにした．

こんな感じのことができます．

def multi_index_dict_test2():
    drawer1 = lambda: MultiIndex(name=Unique(lambda x: x.name, OrderedDict),
                                 sex=NonUnique(lambda x: x.sex))
    drawer2 = lambda: MultiIndex(name=Unique(lambda x: x.name, OrderedDict),
                                 age=NonUnique(lambda x: x.age))
    mdict = MultiIndex(name=Unique(lambda x: x.name, OrderedDict),
                       age=NonUnique(lambda x: x.age, drawer=drawer1),
                       sex=NonUnique(lambda x: x.sex, drawer=drawer2))

    mdict.append(Employee('Mike', 20, 'male'))
    mdict.append(Employee('Jane', 21, 'female'))
    mdict.append(Employee('Bob', 21, 'male'))

    print(mdict['name']['Mike'])
    print(mdict['age'][21]['sex']['female'])  # age == 21 && sex == 'female' をさがす

まとめ

boost::MultiIndex にインスパイアされて Python MultiIndex を実装しました．

全部 hash map でいいじゃん? ということで全部 dict を使って実装できました．

Unique Index に同じものがきたときどうするか，とか考えなきゃいけないことがいっぱいありますが，だいたいはできたと思います．

あと，ぜったいに mdict['name']['Mike'].set_name('マイク') みたいなことはしちゃだめです．

Index にしていいのは Unmutable なものだけですよ！