Add functools.lfu_cache() and functools.lru_cache().

2 files changed, 140 insertions, 2 deletions

diff --git a/Lib/functools.py b/Lib/functools.py
index 1a1f22e..863706d 100644
--- a/Lib/functools.py
+++ b/Lib/functools.py
@@ -4,10 +4,17 @@
 # to allow utilities written in Python to be added
 # to the functools module.
 # Written by Nick Coghlan <ncoghlan at gmail.com>
-#   Copyright (C) 2006 Python Software Foundation.
+# and Raymond Hettinger <python at rcn.com>
+#   Copyright (C) 2006-2010 Python Software Foundation.
 # See C source code for _functools credits/copyright
 
+__all__ = ['update_wrapper', 'wraps', 'WRAPPER_ASSIGNMENTS', 'WRAPPER_UPDATES',
+           'total_ordering', 'cmp_to_key', 'lfu_cache', 'lru_cache']
+
 from _functools import partial, reduce
+from collections import OrderedDict, Counter
+from heapq import nsmallest
+from operator import itemgetter
 
 # update_wrapper() and wraps() are tools to help write
 # wrapper functions that can handle naive introspection
@@ -97,3 +104,88 @@ def cmp_to_key(mycmp):
         def __hash__(self):
             raise TypeError('hash not implemented')
     return K
+
+def lfu_cache(maxsize=100):
+    '''Least-frequently-used cache decorator.
+
+    Arguments to the cached function must be hashable.
+    Cache performance statistics stored in f.hits and f.misses.
+    Clear the cache using f.clear().
+    http://en.wikipedia.org/wiki/Cache_algorithms#Least-Frequently_Used
+
+    '''
+    def decorating_function(user_function):
+        cache = {}                      # mapping of args to results
+        use_count = Counter()           # times each key has been accessed
+        kwd_mark = object()             # separate positional and keyword args
+
+        @wraps(user_function)
+        def wrapper(*args, **kwds):
+            key = args
+            if kwds:
+                key += (kwd_mark,) + tuple(sorted(kwds.items()))
+            use_count[key] += 1         # count a use of this key
+            try:
+                result = cache[key]
+                wrapper.hits += 1
+            except KeyError:
+                result = user_function(*args, **kwds)
+                cache[key] = result
+                wrapper.misses += 1
+                if len(cache) > maxsize:
+                    # purge the 10% least frequently used entries
+                    for key, _ in nsmallest(maxsize // 10,
+                                            use_count.items(),
+                                            key=itemgetter(1)):
+                        del cache[key], use_count[key]
+            return result
+
+        def clear():
+            'Clear the cache and cache statistics'
+            cache.clear()
+            use_count.clear()
+            wrapper.hits = wrapper.misses = 0
+
+        wrapper.hits = wrapper.misses = 0
+        wrapper.clear = clear
+        return wrapper
+    return decorating_function
+
+def lru_cache(maxsize=100):
+    '''Least-recently-used cache decorator.
+
+    Arguments to the cached function must be hashable.
+    Cache performance statistics stored in f.hits and f.misses.
+    Clear the cache using f.clear().
+    http://en.wikipedia.org/wiki/Cache_algorithms#Least_Recently_Used
+
+    '''
+    def decorating_function(user_function):
+        cache = OrderedDict()           # ordered least recent to most recent
+        kwd_mark = object()             # separate positional and keyword args
+
+        @wraps(user_function)
+        def wrapper(*args, **kwds):
+            key = args
+            if kwds:
+                key += (kwd_mark,) + tuple(sorted(kwds.items()))
+            try:
+                result = cache.pop(key)
+                wrapper.hits += 1
+            except KeyError:
+                result = user_function(*args, **kwds)
+                wrapper.misses += 1
+                if len(cache) >= maxsize:
+                    cache.popitem(0)    # purge least recently used cache entry
+            cache[key] = result         # record recent use of this key
+            return result
+
+        def clear():
+            'Clear the cache and cache statistics'
+            cache.clear()
+            wrapper.hits = wrapper.misses = 0
+
+        wrapper.hits = wrapper.misses = 0
+        wrapper.clear = clear
+        return wrapper
+    return decorating_function
diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py
index f6ccc87..a02d37c 100644
--- a/Lib/test/test_functools.py
+++ b/Lib/test/test_functools.py
@@ -4,6 +4,7 @@ import unittest
 from test import support
 from weakref import proxy
 import pickle
+from random import choice
 
 @staticmethod
 def PythonPartial(func, *args, **keywords):
@@ -454,6 +455,50 @@ class TestTotalOrdering(unittest.TestCase):
             class A:
                 pass
 
+class TestLRU(unittest.TestCase):
+
+    def test_lru(self):
+        def orig(x, y):
+            return 3*x+y
+        f = functools.lru_cache(maxsize=20)(orig)
+
+        domain = range(5)
+        for i in range(1000):
+            x, y = choice(domain), choice(domain)
+            actual = f(x, y)
+            expected = orig(x, y)
+            self.assertEquals(actual, expected)
+        self.assert_(f.hits > f.misses)
+        self.assertEquals(f.hits + f.misses, 1000)
+
+        f.clear()   # test clearing
+        self.assertEqual(f.hits, 0)
+        self.assertEqual(f.misses, 0)
+        f(x, y)
+        self.assertEqual(f.hits, 0)
+        self.assertEqual(f.misses, 1)
+
+    def test_lfu(self):
+        def orig(x, y):
+            return 3*x+y
+        f = functools.lfu_cache(maxsize=20)(orig)
+
+        domain = range(5)
+        for i in range(1000):
+            x, y = choice(domain), choice(domain)
+            actual = f(x, y)
+            expected = orig(x, y)
+            self.assertEquals(actual, expected)
+        self.assert_(f.hits > f.misses)
+        self.assertEquals(f.hits + f.misses, 1000)
+
+        f.clear()   # test clearing
+        self.assertEqual(f.hits, 0)
+        self.assertEqual(f.misses, 0)
+        f(x, y)
+        self.assertEqual(f.hits, 0)
+        self.assertEqual(f.misses, 1)
+
 def test_main(verbose=None):
     test_classes = (
         TestPartial,
@@ -461,7 +506,8 @@ def test_main(verbose=None):
         TestPythonPartial,
         TestUpdateWrapper,
         TestWraps,
-        TestReduce
+        TestReduce,
+        TestLRU,
     )
     support.run_unittest(*test_classes)