Skip to content

Commit fefdf29

Browse files
douglas-raillard-armmarcbonnici
authored andcommitted
utils/asyn: Add memoize_method() decorator
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.
1 parent 0ea9c73 commit fefdf29

File tree

2 files changed

+83
-0
lines changed

2 files changed

+83
-0
lines changed

devlib/module/cpufreq.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ def __init__(self, target):
5959
self._governor_tunables = {}
6060

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

295296
@asyn.asyncf
297+
@asyn.memoized_method
296298
async def list_frequencies(self, cpu):
297299
"""Returns a sorted list of frequencies supported by the cpu or an empty list
298300
if not could be found."""

devlib/utils/asyn.py

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
import contextlib
2626
import pathlib
2727
import os.path
28+
import inspect
2829

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

213214

215+
class memoized_method:
216+
"""
217+
Decorator to memmoize a method.
218+
219+
It works for:
220+
221+
* async methods (coroutine functions)
222+
* non-async methods
223+
* method already decorated with :func:`devlib.asyn.asyncf`.
224+
225+
.. note:: This decorator does not rely on hacks to hash unhashable data. If
226+
such input is required, it will either have to be coerced to a hashable
227+
first (e.g. converting a list to a tuple), or the code of
228+
:func:`devlib.asyn.memoized_method` will have to be updated to do so.
229+
"""
230+
def __init__(self, f):
231+
memo = self
232+
233+
sig = inspect.signature(f)
234+
235+
def bind(self, *args, **kwargs):
236+
bound = sig.bind(self, *args, **kwargs)
237+
bound.apply_defaults()
238+
key = (bound.args[1:], tuple(sorted(bound.kwargs.items())))
239+
240+
return (key, bound.args, bound.kwargs)
241+
242+
def get_cache(self):
243+
try:
244+
cache = self.__dict__[memo.name]
245+
except KeyError:
246+
cache = {}
247+
self.__dict__[memo.name] = cache
248+
return cache
249+
250+
251+
if inspect.iscoroutinefunction(f):
252+
@functools.wraps(f)
253+
async def wrapper(self, *args, **kwargs):
254+
cache = get_cache(self)
255+
key, args, kwargs = bind(self, *args, **kwargs)
256+
try:
257+
return cache[key]
258+
except KeyError:
259+
x = await f(*args, **kwargs)
260+
cache[key] = x
261+
return x
262+
else:
263+
@functools.wraps(f)
264+
def wrapper(self, *args, **kwargs):
265+
cache = get_cache(self)
266+
key, args, kwargs = bind(self, *args, **kwargs)
267+
try:
268+
return cache[key]
269+
except KeyError:
270+
x = f(*args, **kwargs)
271+
cache[key] = x
272+
return x
273+
274+
275+
self.f = wrapper
276+
self._name = f.__name__
277+
278+
@property
279+
def name(self):
280+
return '__memoization_cache_of_' + self._name
281+
282+
def __call__(self, *args, **kwargs):
283+
return self.f(*args, **kwargs)
284+
285+
def __get__(self, obj, owner=None):
286+
return self.f.__get__(obj, owner)
287+
288+
def __set__(self, obj, value):
289+
raise RuntimeError("Cannot monkey-patch a memoized function")
290+
291+
def __set_name__(self, owner, name):
292+
self.name = name
293+
294+
214295
def asyncf(f):
215296
"""
216297
Decorator used to turn a coroutine into a blocking function, with an

0 commit comments

Comments
 (0)