multiindex.py

#!/usr/bin/python -utt

from sortedcontainers import SortedList, SortedDict
from operator import itemgetter, attrgetter


class IndexedList(object):
    def __init__(self, hash=None, *args, **kwargs):
        self.pk = 0
        self.primary = SortedDict()
        self.indices = {}
        self.uniques = {}
        self.keys = {}
        self.hash = (lambda x: x) if hash is None else hash
        self.load = kwargs.get('load', 1000)

    def add_unique_index(self, name, key=None):
        if key is None:
            key = attrgetter(name)

        key_getter = lambda pair: self.hash(key(pair[1]))

        self.uniques[name] = SortedDict(self.load)
        self.keys[name] = key_getter

    def add_index(self, name, key=None):
        if key is None:
            key = attrgetter(name)

        self.indices[name] = SortedList(key=lambda pair: key(pair[1]), load=self.load)

    def append(self, item):
        pair = (self.pk, item)
        self.primary[self.pk] = item
        self.pk += 1

        for name, key_getter in self.keys.iteritems():
            self.uniques[name][key_getter(pair)] = pair
        for name in self.indices:
            self.indices[name].add(pair)

    def replace(self, key, val, item):
        old_pair = self.uniques[key][self.hash(val)]
        new_pair = (old_pair[0], item)

        for key, key_getter in self.keys.iteritems():
            self.uniques[key][key_getter(old_pair)] = new_pair

        for key in self.indices:
            self.indices[key].remove(old_pair)
            self.indices[key].add(new_pair)

        self.primary[old_pair[0]] = item

    def insert_or_replace(self, key, val, item):
        if self.hash(val) in self.uniques[key]:
            self.replace(key, val, item)
        else:
            self.append(item)

    def pop_all(self):
        result = []
        for position in self.primary:
            result.append(self.pop(position))

        return result

    def pop(self, position):
        item = self.primary[position]
        del self.primary[position]
        for index, key_getter in self.keys.iteritems():
            del self.uniques[index][key_getter((position, item))]
        for index in self.indices:
            self.indices[index].remove((position, item))
        return item

    def remove(self, position):
        self.pop(position)

    def pop_item(self, item):
        pk = None
        for key, key_getter in self.keys.items():
            if not pk:
                pair = self.uniques[key][key_getter((None, item))]
                pk = pair[0]
                return self.pop(pk)

    def remove_item(self, item):
        self.pop_item(item)

    def remove_unique(self, index, key):
        self.pop_unique(index, key)

    def pop_unique(self, index, key):
        item = self.uniques[index][self.hash(key)]
        return self.pop(item[0])

    def remove_slice(self, index, begin, end, include_last=True):
        self.pop_slice(index, begin, end, include_last)

    def pop_outnumbers(self, index, keep, keep_from_beginning=False):
        """
        Remove all elements but given `keep` amount from the end.
        :param index: index name to sort by
        :param keep: amount of elements to leave
        :param keep_from_beginning: True if we need to leave `keep` first elements instead of `keep` last ones
        """
        target = self.indices[index][keep:] if keep_from_beginning else self.indices[index][:-keep]
        positions = map(itemgetter(0), target)
        removed = map(self.pop, positions)
        return removed

    def pop_slice(self, index, begin, end, include_last=True, limit=0):
        """
        Remove slice of elements by non-unique key
        :param index: index name to slice
        :param begin: elements (>= begin) will be removed
        :param end: elements (<[=] end) will be removed
        :param include_last: inclusive last end
        :param report: False if should not return list of removed objects
        :return: removed object in list
        """
        items = map(itemgetter(0), self.get_slice_items(index, begin, end, include_last, limit))
        removed = map(self.pop, items)
        return removed

    def pop_front(self, index):
        element = self.indices[index][0]
        self.remove(element[0])
        return element

    def __getitem__(self, item):
        return self.primary[item][1]

    def __iter__(self):
        return self.primary.itervalues()

    def __len__(self):
        return len(self.primary)

    def get_max_pk(self):
        return self.pk

    def get_by(self, index, key):
        if index in self.uniques:
            return self.uniques[index][self.hash(key)][-1]
        elif index in self.indices:
            return self.indices[index][key][-1]
        else:
            raise KeyError("{} isn't indexed".format(index))

    def get_slice_items(self, index, begin, end, include_last=True, limit=0):
        start = self.indices[index].bisect_key_left(begin)
        finish = self.indices[index].bisect_key_right(end) if include_last else self.indices[index].bisect_key_left(end)
        return self.indices[index][start:finish if limit == 0 else min(finish, start+limit)]

    def get_slice(self, index, begin, end, include_last=True):
        """
        Get slice of elements by non-unique key
        :param index: index name to slice
        :param begin: elements (>= begin) will be removed
        :param end: elements (<[=] end) will be removed
        :param include_last: inclusive last end
        :return: slice in list
        """
        return map(itemgetter(1), self.get_slice_items(index, begin, end, include_last))

    def try_get_by(self, index, key, default=None):
        try:
            return self.get_by(index, key)
        except KeyError:
            return default

    def iter_by(self, index):
        if index in self.uniques:
            target = self.uniques[index]
        elif index in self.indices:
            target = self.indices[index]
        else:
            raise KeyError("{} isn't indexed".format(index))

        def gen():
            it = target.itervalues()
            while True:
                el = it.next()
                if not isinstance(el, tuple):
                    raise StopIteration
                yield el[1]

        return gen()

    def iter_reverse_by(self, index):
        if index in self.uniques:
            target = self.uniques[index]
        elif index in self.indices:
            target = self.indices[index]
        else:
            raise KeyError("{} isn't indexed".format(index))

        def gen():
            it = target.__reversed__()
            while True:
                ukey = it.next()
                el = self.uniques[index][ukey]
                if not isinstance(el, tuple):
                    continue
                yield el[1]

        return gen()

    def iterate(self, index, begin, end):
        start = self.indices[index].bisect_key_left(begin)
        finish = self.indices[index].bisect_key_right(end)

        def gen():
            it = iter(self.indices[index][start:finish])
            while True:
                el = it.next()
                if not isinstance(el, tuple):
                    raise StopIteration
                yield el[1]

        return gen()