From ad2f39d5e38c1ca98458ae7ea8f37a563f84d772 Mon Sep 17 00:00:00 2001 From: Lucas Moura Date: Thu, 23 Mar 2017 15:30:12 -0300 Subject: [PATCH] Use lru_cache decorator on memoize function --- utils.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/utils.py b/utils.py index 73dd63d63..7a547c67c 100644 --- a/utils.py +++ b/utils.py @@ -270,13 +270,10 @@ def isclose(a, b, rel_tol=1e-09, abs_tol=0.0): # Misc Functions -# TODO: Use functools.lru_cache memoization decorator - - -def memoize(fn, slot=None): +def memoize(fn, slot=None, maxsize=32): """Memoize fn: make it remember the computed value for any argument list. If slot is specified, store result in that slot of first argument. - If slot is false, store results in a dictionary.""" + If slot is false, use lru_cache for caching the values.""" if slot: def memoized_fn(obj, *args): if hasattr(obj, slot): @@ -286,12 +283,9 @@ def memoized_fn(obj, *args): setattr(obj, slot, val) return val else: + @functools.lru_cache(maxsize=maxsize) def memoized_fn(*args): - if args not in memoized_fn.cache: - memoized_fn.cache[args] = fn(*args) - return memoized_fn.cache[args] - - memoized_fn.cache = {} + return fn(*args) return memoized_fn