diff --git a/pytools/__init__.py b/pytools/__init__.py index 688c9cd9..b68002f9 100644 --- a/pytools/__init__.py +++ b/pytools/__init__.py @@ -666,30 +666,34 @@ def wrapper(*args): FunctionValueCache = memoize -class _HasKwargs: - pass - - def memoize_on_first_arg(function, cache_dict_name=None): """Like :func:`memoize_method`, but for functions that take the object in which do memoization information is stored as first argument. Supports cache deletion via ``function_name.clear_cache(self)``. + + .. note:: + + If *function* is nested in another function, a separate cache + is created for each invocation of the surrounding function, + and use of the function in the cache key may keep references + to the nested function alive longer than desired. """ if cache_dict_name is None: - cache_dict_name = intern( - f"_memoize_dic_{function.__module__}{function.__name__}" - ) + cache_dict_name = (memoize_on_first_arg, function) + else: + if not isinstance(cache_dict_name, tuple): + cache_dict_name = (cache_dict_name,) def wrapper(obj, *args, **kwargs): if kwargs: - key = (_HasKwargs, frozenset(kwargs.items())) + args + key = cache_dict_name + args + (frozenset(kwargs.items()),) else: - key = args + key = (cache_dict_name + args) try: - return getattr(obj, cache_dict_name)[key] + return obj._memoize_on_first_arg_dict[key] except AttributeError: attribute_error = True except KeyError: @@ -697,14 +701,14 @@ def wrapper(obj, *args, **kwargs): result = function(obj, *args, **kwargs) if attribute_error: - object.__setattr__(obj, cache_dict_name, {key: result}) + object.__setattr__(obj, "_memoize_on_first_arg_dict", {key: result}) return result else: - getattr(obj, cache_dict_name)[key] = result + obj._memoize_on_first_arg_dict[key] = result return result def clear_cache(obj): - object.__delattr__(obj, cache_dict_name) + object.__delattr__(obj, "_memoize_on_first_arg_dict") from functools import update_wrapper new_wrapper = update_wrapper(wrapper, function) @@ -722,13 +726,12 @@ def memoize_method(method: F) -> F: (e.g. by overwritting ``__setattr__``), e.g. frozen :mod:`dataclasses`. """ - return memoize_on_first_arg(method, - intern(f"_memoize_dic_{method.__name__}")) + return memoize_on_first_arg(method, (memoize_method, method)) class keyed_memoize_on_first_arg: # noqa: N801 - """Like :func:`memoize_method`, but for functions that take the object - in which memoization information is stored as first argument. + """Like :func:`memoize_on_first_arg`, but uses a user-supplied function + for cache key computation. Supports cache deletion via ``function_name.clear_cache(self)``. @@ -738,6 +741,13 @@ class keyed_memoize_on_first_arg: # noqa: N801 used to hold the cache. .. versionadded :: 2020.3 + + .. note:: + + If *function* is nested in another function, a separate cache + is created for each invocation of the surrounding function, + and use of the function in the cache key may keep references + to the nested function alive longer than desired. """ def __init__(self, key, cache_dict_name=None): @@ -798,6 +808,10 @@ def _default_cache_dict_name(self, function): return intern(f"_memoize_dic_{function.__name__}") +class _HasKwargs: + pass + + def memoize_method_with_uncached(uncached_args=None, uncached_kwargs=None): """Supports cache deletion via ``method_name.clear_cache(self)``.