Skip to content

Commit

Permalink
utils/asyn: Add memoize_method() decorator
Browse files Browse the repository at this point in the history
Add a memoize_method decorator that works for async methods. It will not
leak memory since the memoization cache is held in the instance
__dict__, and it does not rely on hacks to hash unhashable data.
  • Loading branch information
douglas-raillard-arm authored and marcbonnici committed Jul 28, 2022
1 parent 0ea9c73 commit fefdf29
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 0 deletions.
2 changes: 2 additions & 0 deletions devlib/module/cpufreq.py
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ def __init__(self, target):
self._governor_tunables = {}

@asyn.asyncf
@asyn.memoized_method
async def list_governors(self, cpu):
"""Returns a list of governors supported by the cpu."""
if isinstance(cpu, int):
Expand Down Expand Up @@ -293,6 +294,7 @@ async def set_governor_tunables(self, cpu, governor=None, per_cpu=None, **kwargs
raise TargetStableError(message)

@asyn.asyncf
@asyn.memoized_method
async def list_frequencies(self, cpu):
"""Returns a sorted list of frequencies supported by the cpu or an empty list
if not could be found."""
Expand Down
81 changes: 81 additions & 0 deletions devlib/utils/asyn.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
import contextlib
import pathlib
import os.path
import inspect

# Allow nesting asyncio loops, which is necessary for:
# * Being able to call the blocking variant of a function from an async
Expand Down Expand Up @@ -211,6 +212,86 @@ def __getattr__(self, attr):
return getattr(self.asyn, attr)


class memoized_method:
"""
Decorator to memmoize a method.
It works for:
* async methods (coroutine functions)
* non-async methods
* method already decorated with :func:`devlib.asyn.asyncf`.
.. note:: This decorator does not rely on hacks to hash unhashable data. If
such input is required, it will either have to be coerced to a hashable
first (e.g. converting a list to a tuple), or the code of
:func:`devlib.asyn.memoized_method` will have to be updated to do so.
"""
def __init__(self, f):
memo = self

sig = inspect.signature(f)

def bind(self, *args, **kwargs):
bound = sig.bind(self, *args, **kwargs)
bound.apply_defaults()
key = (bound.args[1:], tuple(sorted(bound.kwargs.items())))

return (key, bound.args, bound.kwargs)

def get_cache(self):
try:
cache = self.__dict__[memo.name]
except KeyError:
cache = {}
self.__dict__[memo.name] = cache
return cache


if inspect.iscoroutinefunction(f):
@functools.wraps(f)
async def wrapper(self, *args, **kwargs):
cache = get_cache(self)
key, args, kwargs = bind(self, *args, **kwargs)
try:
return cache[key]
except KeyError:
x = await f(*args, **kwargs)
cache[key] = x
return x
else:
@functools.wraps(f)
def wrapper(self, *args, **kwargs):
cache = get_cache(self)
key, args, kwargs = bind(self, *args, **kwargs)
try:
return cache[key]
except KeyError:
x = f(*args, **kwargs)
cache[key] = x
return x


self.f = wrapper
self._name = f.__name__

@property
def name(self):
return '__memoization_cache_of_' + self._name

def __call__(self, *args, **kwargs):
return self.f(*args, **kwargs)

def __get__(self, obj, owner=None):
return self.f.__get__(obj, owner)

def __set__(self, obj, value):
raise RuntimeError("Cannot monkey-patch a memoized function")

def __set_name__(self, owner, name):
self.name = name


def asyncf(f):
"""
Decorator used to turn a coroutine into a blocking function, with an
Expand Down

0 comments on commit fefdf29

Please sign in to comment.