source: branches/decorate/memoize.py @ 855

#!/usr/bin/env python
#
# Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
# Copyright (c) 2012-2016 California Institute of Technology.
# License: 3-clause BSD.  The full license text is available at:
#  - http://mmckerns.github.io/project/mystic/browser/mystic/LICENSE
"""
decorators that cache results to memory, to file, or to a database
"""
from functools import update_wrapper, wraps as functools_wraps
from klepto.rounding import deep_round, simple_round
from klepto.archives import cache as archive_dict
from klepto.keymaps import stringmap
from klepto import inf_cache
from klepto.tools import CacheInfo

__all__ = ['memoize', 'memoized']

#XXX: use cache maxsize algorithms... where dump to archive if maxsize
#XXX: memoized fails when decorating a class method ???

# backward compatibility
memoized = inf_cache

class memoize(object):
    """This is safecache.inf_cache implemented as a class-based decorator.

    Decorator that memoizes a function's return value each time it is called.
    If called later with the same arguments, the memoized value is returned, and
    not re-evaluated.  This may lead to memory issues, as cache is not cleared.
    Can memoize a *method* on an object.
    """
    def __init__(self, cache=None, keymap=None, tol=None, deep=False):
        self.__maxsize = None
        if keymap is None: keymap = stringmap(flat=False)
        if cache is None: cache = archive_dict()
        elif type(cache) is dict: cache = archive_dict(cache)
        self.__cache = cache
        self.__keymap = keymap

        if deep: rounded = deep_round
        else: rounded = simple_round

        @rounded(tol)
        def rounded_args(*args, **kwds):
            return (args, kwds)
        self.__rounded_args = rounded_args
        return

    def __call__(self, user_function):  # i.e. 'decorating_function'
        stats = [0, 0, 0]               # make statistics updateable non-locally
        HIT, MISS, LOAD = 0, 1, 2       # names for the stats fields

        def wrapper(*args, **kwds):
            try:
                _args, _kwds = self.__rounded_args(*args, **kwds)
                key = self.__keymap(*_args, **_kwds)
            except: #TypeError
                result = user_function(*args, **kwds)
                return result

            try:
                # get cache entry
                result = self.__cache[key]
                stats[HIT] += 1
            except KeyError:
                # if not in cache, look in archive
                if self.__cache.archived():
                    self.__cache.load(key)
                try:
                    result = self.__cache[key]
                    stats[LOAD] += 1
                except KeyError:
                    # if not found, then compute
                    result = user_function(*args, **kwds)
                    self.__cache[key] = result
                    stats[MISS] += 1
            return result

        def archive(obj):
            """Replace the cache archive"""
            self.__cache.archive = obj

        def __get_cache():
            """Get the cache"""
            return self.__cache

        def clear(keepstats=False):
            """Clear the cache and statistics"""
            self.__cache.clear()
            if not keepstats: stats[:] = [0, 0, 0]

        def info():
            """Report cache statistics"""
            return CacheInfo(stats[HIT], stats[MISS], stats[LOAD], self.__maxsize, len(self.__cache))

        # interface
        wrapper.__wrapped__ = user_function
        wrapper.info = info
        wrapper.clear = clear
        wrapper.load = self.__cache.load
        wrapper.dump = self.__cache.dump
        wrapper.archive = archive
        wrapper.archived = self.__cache.archived
        wrapper.__cache__ = __get_cache
        return update_wrapper(wrapper, user_function)

#   def __repr__(self):
#       """Return the function's docstring."""
#       return self.func.__doc__

    def __get__(self, obj, objtype):
        """Support instance methods."""
        import functools
        return functools.partial(self.__call__, obj)

# EOF