# HG changeset patch
# User Radomir Dopieralski <devel@sheep.art.pl>
# Date 1386082968 -3600
# Node ID 66c1c9f3256785552224ca177ed77a8312de6bb8
# Parent  373797990f57566dc9fe1aab28ba2e92a89c4e99
Add use_weakref option to functools.lru_cache

The use_weakref option makes lru_cache create weak references to the function
arguments when creating the cache key, which prevents them from staying in the
memory forever even after all other references to them are gone.

This is a proof of concept for http://bugs.python.org/issue19859

diff -r 373797990f57 -r 66c1c9f32567 Lib/functools.py
--- a/Lib/functools.py	Mon Dec 02 02:56:02 2013 +0100
+++ b/Lib/functools.py	Tue Dec 03 16:02:48 2013 +0100
@@ -338,7 +338,14 @@
     def __hash__(self):
         return self.hashvalue
 
-def _make_key(args, kwds, typed,
+def _weaken(arg, callback=None):
+    """Try to make weak reference. Return the original object if it fails."""
+    try:
+        return weakref.ref(arg, callback)
+    except TypeError:
+        return arg
+
+def _make_key(args, kwds, typed, remove_key=None,
              kwd_mark = (object(),),
              fasttypes = {int, str, frozenset, type(None)},
              sorted=sorted, tuple=tuple, type=type, len=len):
@@ -366,7 +373,41 @@
         return key[0]
     return _HashedSeq(key)
 
-def lru_cache(maxsize=128, typed=False):
+def _make_weakref_key(args, kwds, typed, remove_key,
+                      kwd_mark = (object(),),
+                      fasttypes= {int, str, frozenset, type(None)},
+                      sorted=sorted, tuple=tuple, type=type, len=len):
+    """Like ``_make_key``, but use weak references to arguments.
+
+    The function ``remove_key`` will be called back with a key when any of that
+    key's arguments no longer has any references.
+
+    This function duplicates some code from ``_make_key`` for performance
+    reasons, to avoid additional checks and function calls in the original
+    function.
+    """
+    key = None
+
+    def weakref_callback(ref):
+        # The key here is stored in a closure, because we need this callback
+        # before we have the key, as it is used in the weakrefs for the key.
+        remove_key(key)
+
+    key = tuple(_weaken(arg, weakref_callback) for arg in args)
+    if kwds:
+        sorted_items = sorted(kwds.items())
+        key += kwd_mark
+        for key, value in sorted_items:
+            key += key, _weaken(value, weakref_callback)
+    if typed:
+        key += tuple(type(v) for v in args)
+        if kwds:
+            key += tuple(type(v) for k, v in sorted_items)
+    elif len(key) == 1 and type(key[0]) in fasttypes:
+        return key[0]
+    return _HashedSeq(key)
+
+def lru_cache(maxsize=128, typed=False, use_weakref=False):
     """Least-recently-used cache decorator.
 
     If *maxsize* is set to None, the LRU features are disabled and the cache
@@ -376,6 +417,11 @@
     For example, f(3.0) and f(3) will be treated as distinct calls with
     distinct results.
 
+    If *use_weakref* is True, the cache will not keep references to the
+    arguments, keeping them alive in memory forever. Instead, cache entries
+    related to objects that are no longer used will be automatically removed
+    from cache.
+
     Arguments to the cached function must be hashable.
 
     View the cache statistics named tuple (hits, misses, maxsize, currsize)
@@ -393,9 +439,13 @@
 
     # Constants shared by all lru cache instances:
     sentinel = object()          # unique object used to signal cache misses
-    make_key = _make_key         # build a key from the function arguments
     PREV, NEXT, KEY, RESULT = 0, 1, 2, 3   # names for the link fields
 
+    if use_weakref:
+        make_key = _make_weakref_key
+    else:
+        make_key = _make_key         # build a key from the function arguments
+
     def decorating_function(user_function):
         cache = {}
         hits = misses = 0
@@ -404,6 +454,8 @@
         lock = RLock()           # because linkedlist updates aren't threadsafe
         root = []                # root of the circular doubly linked list
         root[:] = [root, root, None, None]     # initialize by pointing to self
+        remove_key = None        # only used with use_weakref
+
 
         if maxsize == 0:
 
@@ -415,11 +467,19 @@
                 return result
 
         elif maxsize is None:
+            if use_weakref:
+
+                def remove_key(key):
+                    # A callback for removing keys from cache
+                    try:
+                        del cache[key]
+                    except KeyError:
+                        pass
 
             def wrapper(*args, **kwds):
                 # Simple caching without ordering or size limit
                 nonlocal hits, misses
-                key = make_key(args, kwds, typed)
+                key = make_key(args, kwds, typed, remove_key)
                 result = cache_get(key, sentinel)
                 if result is not sentinel:
                     hits += 1
@@ -430,11 +490,22 @@
                 return result
 
         else:
+            if use_weakref:
+
+                def remove_key(key):
+                    # A callback for removing keys from cache
+                    with lock:
+                        link = cache_get(key)
+                        if link is not None:
+                            link_prev, link_next, _key, result = link
+                            link_prev[NEXT] = link_next
+                            link_next[PREV] = link_prev
+                            del cache[key]
 
             def wrapper(*args, **kwds):
                 # Size limited caching that tracks accesses by recency
                 nonlocal root, hits, misses, full
-                key = make_key(args, kwds, typed)
+                key = make_key(args, kwds, typed, remove_key)
                 with lock:
                     link = cache_get(key)
                     if link is not None: