nilmdb/nilmdb/utils/lrucache.py

# Memoize a function's return value with a least-recently-used cache
# Based on:
#   http://code.activestate.com/recipes/498245-lru-and-lfu-cache-decorators/
# with added 'destructor' functionality.

import collections
import decorator

def lru_cache(size = 10, onremove = None, keys = slice(None)):
    """Least-recently-used cache decorator.

    @lru_cache(size = 10, onevict = None)
    def f(...):
        pass

    Given a function and arguments, memoize its return value.  Up to
    'size' elements are cached.  'keys' is a slice object that
    represents which arguments are used as the cache key.

    When evicting a value from the cache, call the function
    'onremove' with the value that's being evicted.

    Call f.cache_remove(...) to evict the cache entry with the given
    arguments.  Call f.cache_remove_all() to evict all entries.
    f.cache_hits and f.cache_misses give statistics.
    """

    def decorate(func):
        cache = collections.OrderedDict()	# order: least- to most-recent

        def evict(value):
            if onremove:
                onremove(value)

        def wrapper(orig, *args, **kwargs):
            if kwargs:
                raise NotImplementedError("kwargs not supported")
            key = args[keys]
            try:
                value = cache.pop(key)
                orig.cache_hits += 1
            except KeyError:
                value = orig(*args)
                orig.cache_misses += 1
                if len(cache) >= size:
                    evict(cache.popitem(0)[1])	# evict LRU cache entry
            cache[key] = value              	# (re-)insert this key at end
            return value

        def cache_remove(*args):
            """Remove the described key from this cache, if present."""
            key = args
            if key in cache:
                evict(cache.pop(key))
            else:
                if len(cache) > 0 and len(args) != len(cache.iterkeys().next()):
                    raise KeyError("trying to remove from LRU cache, but "
                                   "number of arguments doesn't match the "
                                   "cache key length")

        def cache_remove_all():
            for key in cache:
                evict(cache.pop(key))

        def cache_info():
            return (func.cache_hits, func.cache_misses)

        new = decorator.decorator(wrapper, func)
        func.cache_hits = 0
        func.cache_misses = 0
        new.cache_info = cache_info
        new.cache_remove = cache_remove
        new.cache_remove_all = cache_remove_all
        return new

    return decorate