bpo-40504: Allow weakrefs to lru_cache objects (GH-19938)

3 files changed, 38 insertions, 1 deletions

diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py
index 9503f40..b3893a1 100644
--- a/Lib/test/test_functools.py
+++ b/Lib/test/test_functools.py
@@ -14,6 +14,8 @@ import typing
 import unittest
 import unittest.mock
 import os
+import weakref
+import gc
 from weakref import proxy
 import contextlib
 
@@ -1938,6 +1940,35 @@ class TestLRU:
             return 1
         self.assertEqual(f.cache_parameters(), {'maxsize': 1000, "typed": True})
 
+    def test_lru_cache_weakrefable(self):
+        @self.module.lru_cache
+        def test_function(x):
+            return x
+
+        class A:
+            @self.module.lru_cache
+            def test_method(self, x):
+                return (self, x)
+
+            @staticmethod
+            @self.module.lru_cache
+            def test_staticmethod(x):
+                return (self, x)
+
+        refs = [weakref.ref(test_function),
+                weakref.ref(A.test_method),
+                weakref.ref(A.test_staticmethod)]
+
+        for ref in refs:
+            self.assertIsNotNone(ref())
+
+        del A
+        del test_function
+        gc.collect()
+
+        for ref in refs:
+            self.assertIsNone(ref())
+
 
 @py_functools.lru_cache()
 def py_cached_func(x, y):
diff --git a/Misc/NEWS.d/next/Library/2020-05-05-17-12-47.bpo-40504.EX6wPn.rst b/Misc/NEWS.d/next/Library/2020-05-05-17-12-47.bpo-40504.EX6wPn.rst
new file mode 100644
index 0000000..261a49e
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2020-05-05-17-12-47.bpo-40504.EX6wPn.rst
@@ -0,0 +1 @@
+:func:`functools.lru_cache` objects can now be the targets of weakrefs.
\ No newline at end of file
diff --git a/Modules/_functoolsmodule.c b/Modules/_functoolsmodule.c
index fd4b4c2..d158d3b 100644
--- a/Modules/_functoolsmodule.c
+++ b/Modules/_functoolsmodule.c
@@ -783,6 +783,7 @@ typedef struct lru_cache_object {
     Py_ssize_t misses;
     PyObject *cache_info_type;
     PyObject *dict;
+    PyObject *weakreflist;
 } lru_cache_object;
 
 static PyTypeObject lru_cache_type;
@@ -1196,6 +1197,7 @@ lru_cache_new(PyTypeObject *type, PyObject *args, PyObject *kw)
     Py_INCREF(cache_info_type);
     obj->cache_info_type = cache_info_type;
     obj->dict = NULL;
+    obj->weakreflist = NULL;
     return (PyObject *)obj;
 }
 
@@ -1227,6 +1229,8 @@ lru_cache_dealloc(lru_cache_object *obj)
     lru_list_elem *list;
     /* bpo-31095: UnTrack is needed before calling any callbacks */
     PyObject_GC_UnTrack(obj);
+    if (obj->weakreflist != NULL)
+        PyObject_ClearWeakRefs((PyObject*)obj);
 
     list = lru_cache_unlink_list(obj);
     Py_XDECREF(obj->cache);
@@ -1384,7 +1388,8 @@ static PyTypeObject lru_cache_type = {
     (traverseproc)lru_cache_tp_traverse,/* tp_traverse */
     (inquiry)lru_cache_tp_clear,        /* tp_clear */
     0,                                  /* tp_richcompare */
-    0,                                  /* tp_weaklistoffset */
+    offsetof(lru_cache_object, weakreflist),
+                                        /* tp_weaklistoffset */
     0,                                  /* tp_iter */
     0,                                  /* tp_iternext */
     lru_cache_methods,                  /* tp_methods */