Better `@memoized` decorator.

1 files changed, 25 insertions, 8 deletions

diff --git a/lib/spack/llnl/util/lang.py b/lib/spack/llnl/util/lang.py
index 9e1bef18ca..be6dad867e 100644
--- a/lib/spack/llnl/util/lang.py
+++ b/lib/spack/llnl/util/lang.py
@@ -26,6 +26,7 @@ import os
 import re
 import sys
 import functools
+import collections
 import inspect
 
 # Ignore emacs backups when listing modules
@@ -170,16 +171,32 @@ def has_method(cls, name):
     return False
 
 
-def memoized(obj):
+class memoized(object):
     """Decorator that caches the results of a function, storing them
        in an attribute of that function."""
-    cache = obj.cache = {}
-    @functools.wraps(obj)
-    def memoizer(*args, **kwargs):
-        if args not in cache:
-            cache[args] = obj(*args, **kwargs)
-        return cache[args]
-    return memoizer
+    def __init__(self, func):
+        self.func = func
+        self.cache = {}
+
+
+    def __call__(self, *args):
+        if not isinstance(args, collections.Hashable):
+            # Not hashable, so just call the function.
+            return self.func(*args)
+
+        if args not in self.cache:
+            self.cache[args] = self.func(*args)
+        return self.cache[args]
+
+
+    def __get__(self, obj, objtype):
+        """Support instance methods."""
+        return functools.partial(self.__call__, obj)
+
+
+    def clear(self):
+        """Expunge cache so that self.func will be called again."""
+        self.cache.clear()
 
 
 def list_modules(directory, **kwargs):