Various python dictionary-like data structures.

dictionaries.py

# -*- coding: utf-8 -*-


from collections import defaultdict


class keydefaultdict(defaultdict):
    """Defaultdict that takes inserted key as an argument.

    Example::

        d = keydefaultdict(C)
        d[x]  # returns C(x)

    """

    def __missing__(self, key):
        if self.default_factory is None:
            raise KeyError(key)
        else:
            args = key if isinstance(key, tuple) else (key,)
            ret = self[key] = self.default_factory(*args)
            return ret


class intervaldict(dict):
    """Dictionary that expects intervals of length two as keys.
    Items from the dictionary are retrieved according to the interval
    value.

    Example::

        d = intervaldict({(1, 2): 'a', (8, 10): 'b'})
        d[8.2]  # returns 'b'
        d[1]  # returns 'a'

    """

    def get_interval(self, key):
        if isinstance(key, tuple):
            return key
        for k, v in self.items():
            if k[0] <= key < k[1]:
                return k
        raise KeyError('Key {!r} is not between any values in '
                       'the intervaldict'.format(key))

    def __getitem__(self, key):
        interval = self.get_interval(key)
        return dict.__getitem__(self, interval)

    def __setitem__(self, key, value):
        interval = self.get_interval(key)
        dict.__setitem__(self, interval, value)

    def __contains__(self, key):
        try:
            return bool(self[key]) or True
        except KeyError:
            return False