diff --git a/pytools/__init__.py b/pytools/__init__.py index d5b8d8e9feb525a61ef5fdabeb10a671cd44b211..c7bb9d59dcb45188656cfd0e3e28a71925a1fefe 100644 --- a/pytools/__init__.py +++ b/pytools/__init__.py @@ -74,6 +74,7 @@ Memoization .. autofunction:: memoize_in .. autofunction:: keyed_memoize_on_first_arg .. autofunction:: keyed_memoize_method +.. autofunction:: keyed_memoize_in Argmin/max ---------- @@ -914,6 +915,41 @@ class memoize_in: # noqa return new_inner + +class keyed_memoize_in: # noqa + """Like :class:`memoize_in`, but additionally uses a function *key* to + compute the key under which the function result is memoized. + + :arg key: A function receiving the same arguments as the decorated function + which computes and returns the cache key. + + .. versionadded :: 2021.2.1 + """ + + def __init__(self, container, identifier, key): + try: + memoize_in_dict = container._pytools_keyed_memoize_in_dict + except AttributeError: + memoize_in_dict = {} + object.__setattr__(container, "_pytools_keyed_memoize_in_dict", + memoize_in_dict) + + self.cache_dict = memoize_in_dict.setdefault(identifier, {}) + self.key = key + + def __call__(self, inner): + @wraps(inner) + def new_inner(*args): + key = self.key(*args) + try: + return self.cache_dict[key] + except KeyError: + result = inner(*args) + self.cache_dict[key] = result + return result + + return new_inner + # }}} diff --git a/pytools/version.py b/pytools/version.py index b431cafddb6c25c285c9046eaf4bbf749911478e..fc538bf6b12a3c52bcbb8c65fc0d92819e5c2836 100644 --- a/pytools/version.py +++ b/pytools/version.py @@ -1,3 +1,3 @@ -VERSION = (2021, 2) +VERSION = (2021, 2, 1) VERSION_STATUS = "" VERSION_TEXT = ".".join(str(x) for x in VERSION) + VERSION_STATUS