Bug 1053080 - Improve mozbuild.util.memoize and add memoized_property. r=gps

python/mozbuild/mozbuild/test/test_util.py

@@ -21,6 +21,7 @@ from mozbuild.util import (
     FileAvoidWrite,
     hash_file,
     memoize,
+    memoized_property,
     resolve_target_to_make,
     MozbuildDeletionError,
     HierarchicalStringList,
@@ -445,6 +446,7 @@ class TestMemoize(unittest.TestCase):
             self._count += 1
             return a + b
 
+        self.assertEqual(self._count, 0)
         self.assertEqual(wrapped(1, 1), 2)
         self.assertEqual(self._count, 1)
         self.assertEqual(wrapped(1, 1), 2)
@@ -458,6 +460,53 @@ class TestMemoize(unittest.TestCase):
         self.assertEqual(wrapped(1, 1), 2)
         self.assertEqual(self._count, 3)
 
+    def test_memoize_method(self):
+        class foo(object):
+            def __init__(self):
+                self._count = 0
+
+            @memoize
+            def wrapped(self, a, b):
+                self._count += 1
+                return a + b
+
+        instance = foo()
+        refcount = sys.getrefcount(instance)
+        self.assertEqual(instance._count, 0)
+        self.assertEqual(instance.wrapped(1, 1), 2)
+        self.assertEqual(instance._count, 1)
+        self.assertEqual(instance.wrapped(1, 1), 2)
+        self.assertEqual(instance._count, 1)
+        self.assertEqual(instance.wrapped(2, 1), 3)
+        self.assertEqual(instance._count, 2)
+        self.assertEqual(instance.wrapped(1, 2), 3)
+        self.assertEqual(instance._count, 3)
+        self.assertEqual(instance.wrapped(1, 2), 3)
+        self.assertEqual(instance._count, 3)
+        self.assertEqual(instance.wrapped(1, 1), 2)
+        self.assertEqual(instance._count, 3)
+
+        # Memoization of methods is expected to not keep references to
+        # instances, so the refcount shouldn't have changed after executing the
+        # memoized method.
+        self.assertEqual(refcount, sys.getrefcount(instance))
+
+    def test_memoized_property(self):
+        class foo(object):
+            def __init__(self):
+                self._count = 0
+
+            @memoized_property
+            def wrapped(self):
+                self._count += 1
+                return 42
+
+        instance = foo()
+        self.assertEqual(instance._count, 0)
+        self.assertEqual(instance.wrapped, 42)
+        self.assertEqual(instance._count, 1)
+        self.assertEqual(instance.wrapped, 42)
+        self.assertEqual(instance._count, 1)
 
 if __name__ == '__main__':
     main()

python/mozbuild/mozbuild/util.py

@@ -10,6 +10,7 @@ from __future__ import unicode_literals
 import copy
 import difflib
 import errno
+import functools
 import hashlib
 import os
 import stat
@@ -20,7 +21,6 @@ from collections import (
     defaultdict,
     OrderedDict,
 )
-from functools import wraps
 from StringIO import StringIO
 
 
@@ -723,12 +723,44 @@ class OrderedDefaultDict(OrderedDict):
             return value
 
 
-def memoize(func):
-    cache = {}
+class memoize(dict):
+    '''A decorator to memoize the results of function calls depending
+    on its arguments.
+    Both functions and instance methods are handled, although in the
+    instance method case, the results are cache in the instance itself.
+    '''
+    def __init__(self, func):
+        self.func = func
+        functools.update_wrapper(self, func)
 
-    @wraps(func)
-    def wrapper(*args):
+    def __call__(self, *args):
+        if args not in self:
+            self[args] = self.func(*args)
+        return self[args]
+
+    def method_call(self, instance, *args):
+        name = '_%s' % self.func.__name__
+        if not hasattr(instance, name):
+            setattr(instance, name, {})
+        cache = getattr(instance, name)
         if args not in cache:
-            cache[args] = func(*args)
+            cache[args] = self.func(instance, *args)
         return cache[args]
-    return wrapper
+
+    def __get__(self, instance, cls):
+        return functools.update_wrapper(
+            functools.partial(self.method_call, instance), self.func)
+
+
+class memoized_property(object):
+    '''A specialized version of the memoize decorator that works for
+    class instance properties.
+    '''
+    def __init__(self, func):
+        self.func = func
+
+    def __get__(self, instance, cls):
+        name = '_%s' % self.func.__name__
+        if not hasattr(instance, name):
+            setattr(instance, name, self.func(instance))
+        return getattr(instance, name)