Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Steps towards more symbolic dict keys in memoize_on_first_arg, memoize_method #80

Open
wants to merge 8 commits into
base: main
Choose a base branch
from
48 changes: 31 additions & 17 deletions pytools/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -666,45 +666,49 @@ 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)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this work if function is nested (i.e. not in the global namespace)? It should keep changing its hash/id in that case, right?

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, nested functions are a difficult case, mainly because it's hard to get stable identifiers for them. My understanding is that, at least, it would fail "safe" (separate caches for each function) using this code.

IMO, @memoize_in is more appropriate for nested functions.

Copy link
Contributor

@alexfikl alexfikl Apr 17, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe add a small note to the docs that it's really not recommended for nested?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Besides that, it looks good to me! 👍

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:
attribute_error = False

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)
Expand All @@ -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)``.

Expand All @@ -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):
Expand Down Expand Up @@ -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)``.

Expand Down