lookup_table.py

def lookup_table(values, key=None, keyval=None, unique=False, use_lists=False):
    """
    Builds a dict-based lookup table (index) elegantly.
    Supports building normal and unique lookup tables.  For example:
    >>> lookup_table(['foo', 'bar', 'baz', 'qux', 'quux'],
    ...              lambda s: s[0])
    {
        'b': {'bar', 'baz'},
        'f': {'foo'},
        'q': {'quux', 'qux'}
    }
    For key functions that uniquely identify values, set unique=True:
    >>> lookup_table(['foo', 'bar', 'baz', 'qux', 'quux'],
    ...              lambda s: s[0],
    ...              unique=True)
    {
        'b': 'baz',
        'f': 'foo',
        'q': 'quux'
    }
    The values of the resulting lookup table will be values, not sets.
    For extra power, you can even change the values while building up the LUT.
    To do so, use the `keyval` function instead of the `key` arg:
    >>> lookup_table(['foo', 'bar', 'baz', 'qux', 'quux'],
    ...              keyval=lambda s: (s[0], s[1:]))
    {
        'b': {'ar', 'az'},
        'f': {'oo'},
        'q': {'uux', 'ux'}
    }
    """
    if keyval is None:
        if key is None:
            keyval = lambda v: v
        else:
            keyval = lambda v: (key(v), v)

    if unique:
        return dict(keyval(v) for v in values)

    lut = {}
    for value in values:
        k, v = keyval(value)
        try:
            s = lut[k]
        except KeyError:
            if use_lists:
                s = lut[k] = list()
            else:
                s = lut[k] = set()
        if use_lists:
            s.append(v)
        else:
            s.add(v)
    return dict(lut)


def reverse_lookup_table(lut):
    """
    Takes any dict of the form:
        k => iterable(values)
    And creates a reverse lookup table for it.
    """
    return {v: k for k, values in lut.iteritems()
                 for v in values}  # noqa