cacher.py

#   Copyright © 2017  Ben Longbons
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU Lesser General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.


''' Similar to reprlib + lru_cache, but more general/specialized.
'''

import functools
import inspect


def make_key(args):
    if isinstance(args, dict):
        rv = []
        for k, v in sorted(args.items()):
            v = make_key(v)
            rv.append((k, v))
        return tuple(rv)
    if isinstance(args, set):
        args = sorted(args)
    if isinstance(args, (tuple, list)):
        rv = []
        for a in args:
            a = make_key(a)
            rv.append(a)
        return tuple(rv)
    hash(args) # assert hashability of remaining types
    return args


def use_cache(fn, recursion='<...recursion...>'):
    if not callable(fn):
        return use_cache(fn, recursion=fn)

    signature = inspect.signature(fn)
    default = signature.parameters['cache'].default

    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        _sentinel = object()

        bound_args = signature.bind(*args, **kwargs)
        bound_args.apply_defaults()
        cache = bound_args.arguments.pop('cache')

        if cache is None:
            cache = {}
        key = (fn, make_key(bound_args.arguments))
        rv = cache.get(key, _sentinel)
        if rv is not _sentinel:
            return rv
        cache[key] = recursion
        rv = fn(*args, **kwargs)
        cache[key] = rv
        return rv
    return wrapper