# HG changeset patch # Parent e6a0418ee09a884850abb82e8d0fbeb8aeca67c8 [mq]: functools.lru_cache-in-c diff --git a/Lib/functools.py b/Lib/functools.py --- a/Lib/functools.py +++ b/Lib/functools.py @@ -268,3 +268,16 @@ return wrapper return decorating_function + + +try: + from _functools import _lru_cache +except ImportError: + pass +else: + def lru_cache(*args, **kwds): + def decorating_function(user_function): + wrapper = _lru_cache(user_function, *args, **kwds) + update_wrapper(wrapper=wrapper, wrapped=user_function) + return wrapper + return decorating_function diff --git a/Lib/test/test_functools.py b/Lib/test/test_functools.py --- a/Lib/test/test_functools.py +++ b/Lib/test/test_functools.py @@ -623,6 +623,7 @@ def test_lru(self): def orig(x, y): + '''does a thingy''' return 3*x+y f = functools.lru_cache(maxsize=20)(orig) hits, misses, maxsize, currsize = f.cache_info() @@ -756,6 +757,15 @@ self.assertEqual(square.cache_info().hits, 4) self.assertEqual(square.cache_info().misses, 4) + def test_lru_cache_decoration(self): + def f(zomg: 'zomg_annotation'): + '''f doc string''' + return 42 + g = functools.lru_cache()(f) + for attr in functools.WRAPPER_ASSIGNMENTS: + self.assertEqual(getattr(g, attr), getattr(f, attr)) + + def test_main(verbose=None): test_classes = ( TestPartial, diff --git a/Modules/_functoolsmodule.c b/Modules/_functoolsmodule.c --- a/Modules/_functoolsmodule.c +++ b/Modules/_functoolsmodule.c @@ -540,6 +540,400 @@ of the sequence in the calculation, and serves as a default when the\n\ sequence is empty."); +/* lru_cache object **********************************************************/ + +/* this object is used delimit args and keywords in the cache keys */ +static PyObject *kwd_mark; + +typedef struct lru_list_elem lru_list_elem; + +typedef struct lru_list_elem { + lru_list_elem *prev, *next; + PyObject *key, *result; +} lru_list_elem; + +typedef struct lru_cache_object lru_cache_object; + +typedef PyObject *(*lru_cache_ternaryfunc)(lru_cache_object *, PyObject *, PyObject *); + +typedef struct lru_cache_object { + PyObject_HEAD + Py_ssize_t maxsize; + PyObject *maxsize_O; + PyObject *func; + lru_cache_ternaryfunc wrapper; + PyObject *cache; + Py_ssize_t misses, hits; + lru_list_elem root; + int typed; + PyObject *dict; +} lru_cache_object; + +static PyTypeObject lru_cache_type; + +static PyObject * +lru_cache_make_key(PyObject *args, PyObject *kwds, int typed) +{ + PyObject *key, *sorted_items; + Py_ssize_t key_size, pos, key_pos; + + /* short path, key will match args anyway, which is a tuple */ + if (!typed && !kwds) { + Py_INCREF(args); + return args; + } + + if (kwds) { + assert(PyDict_Size(kwds)); + if (!(sorted_items = PyDict_Items(kwds))) + return NULL; + if (0 > PyList_Sort(sorted_items)) { + Py_DECREF(sorted_items); + return NULL; + } + } else + sorted_items = NULL; + + key_size = PyTuple_GET_SIZE(args); + if (kwds) + key_size += PyList_GET_SIZE(sorted_items); + if (typed) + key_size *= 2; + if (kwds) + key_size++; + + key = PyTuple_New(key_size); + key_pos = 0; + + for (pos = 0; pos < PyTuple_GET_SIZE(args); ++pos) { + PyObject *item = PyTuple_GET_ITEM(args, pos); + Py_INCREF(item); + PyTuple_SET_ITEM(key, key_pos++, item); + } + if (kwds) { + Py_INCREF(kwd_mark); + PyTuple_SET_ITEM(key, key_pos++, kwd_mark); + for (pos = 0; pos < PyList_GET_SIZE(sorted_items); ++pos) { + PyObject *item = PyList_GET_ITEM(sorted_items, pos); + Py_INCREF(item); + PyTuple_SET_ITEM(key, key_pos++, item); + } + } + if (typed) { + for (pos = 0; pos < PyTuple_GET_SIZE(args); ++pos) { + PyObject *item = (PyObject *)Py_TYPE(PyTuple_GET_ITEM(args, pos)); + Py_INCREF(item); + PyTuple_SET_ITEM(key, key_pos++, item); + } + if (kwds) { + for (pos = 0; pos < PyList_GET_SIZE(sorted_items); ++pos) { + PyObject *item = (PyObject *)Py_TYPE(PyTuple_GET_ITEM(PyList_GET_ITEM(sorted_items, pos), 1)); + Py_INCREF(item); + PyTuple_SET_ITEM(key, key_pos++, item); + } + } + } + assert(key_pos == key_size); + if (kwds) + Py_DECREF(sorted_items); + return key; +} + +static PyObject * +uncached_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwds) +{ + PyObject *result = PyObject_Call(self->func, args, kwds); + if (!result) + return NULL; + self->misses++; + return result; +} + +static PyObject * +infinite_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwds) +{ + PyObject *result; + PyObject *key = lru_cache_make_key(args, kwds, self->typed); + if (!key) + return NULL; + result = PyDict_GetItemWithError(self->cache, key); + if (result) { + Py_INCREF(result); + self->hits++; + Py_DECREF(key); + return result; + } + if (PyErr_Occurred()) { + Py_DECREF(key); + return NULL; + } + result = PyObject_Call(self->func, args, kwds); + if (!result) { + Py_DECREF(key); + return NULL; + } + if (PyDict_SetItem(self->cache, key, result) < 0) { + Py_DECREF(result); + Py_DECREF(key); + return NULL; + } + Py_DECREF(key); + self->misses++; + return result; +} + +static void +lru_cache_list_extricate(lru_list_elem *link) +{ + link->prev->next = link->next; + link->next->prev = link->prev; +} + +static void +lru_cache_list_append(lru_list_elem *root, lru_list_elem *link) +{ + lru_list_elem *last = root->prev; + last->next = root->prev = link; + link->prev = last; + link->next = root; +} + +static PyObject * +bounded_lru_cache_wrapper(lru_cache_object *self, PyObject *args, PyObject *kwds) +{ + PyObject *key = lru_cache_make_key(args, kwds, self->typed); + if (!key) + return NULL; + PyObject *value = PyDict_GetItemWithError(self->cache, key); + if (value) { + lru_list_elem *link = PyCapsule_GetPointer(value, NULL); + lru_cache_list_extricate(link); + lru_cache_list_append(&self->root, link); + self->hits++; + Py_DECREF(key); + Py_INCREF(link->result); + return link->result; + } + if (PyErr_Occurred()) { + Py_DECREF(key); + return NULL; + } + PyObject *result = PyObject_Call(self->func, args, kwds); + if (!result) { + Py_DECREF(key); + return NULL; + } + lru_list_elem *link; + if (PyDict_Size(self->cache) == self->maxsize) { + /* extricate the oldest item */ + link = self->root.next; + lru_cache_list_extricate(link); + /* grab its capsule */ + value = PyDict_GetItem(self->cache, link->key); + Py_INCREF(value); + /* remove its key from the cache */ + if (0 > PyDict_DelItem(self->cache, link->key)) + abort(); + /* scrub the result from the link */ + Py_DECREF(link->result); + } else { + link = PyMem_New(lru_list_elem, 1); + value = PyCapsule_New(link, NULL, NULL); + } + lru_cache_list_append(&self->root, link); + link->key = key; + link->result = result; + Py_INCREF(result); + if (0 > PyDict_SetItem(self->cache, key, value)) abort(); + Py_DECREF(key); + Py_DECREF(value); + self->misses++; + return result; +} + +static PyObject * +lru_cache_new(PyTypeObject *type, PyObject *args, PyObject *kw) +{ + PyObject *func; + PyObject *maxsize_O = NULL; + PyObject *typed_O = NULL; + int typed; + lru_cache_object *obj; + Py_ssize_t maxsize = 100; + PyObject *(*wrapper)(lru_cache_object *, PyObject *, PyObject *); + static char *keywords[] = {"func", "maxsize", "typed", NULL}; + + if (!PyArg_ParseTupleAndKeywords(args, kw, "O|OO:lru_cache", keywords, + &func, &maxsize_O, &typed_O)) { + return NULL; + } + + if (!PyCallable_Check(func)) { + PyErr_SetString(PyExc_TypeError, + "the first argument must be callable"); + return NULL; + } + + if (maxsize_O == NULL) { + maxsize_O = PyLong_FromSsize_t(maxsize); + wrapper = bounded_lru_cache_wrapper; + } else if (maxsize_O == Py_None) { + wrapper = infinite_lru_cache_wrapper; + Py_INCREF(maxsize_O); + } else if (PyNumber_Check(maxsize_O)) { + maxsize = PyNumber_AsSsize_t(maxsize_O, PyExc_OverflowError); + if (maxsize == -1 && PyErr_Occurred()) + return NULL; + if (maxsize == 0) + wrapper = uncached_lru_cache_wrapper; + else + wrapper = bounded_lru_cache_wrapper; + Py_INCREF(maxsize_O); + } else { + PyErr_SetString(PyExc_TypeError, "maxsize should be integer or None"); + return NULL; + } + + if (typed_O) { + int err = PyObject_IsTrue(typed_O); + if (err < 0) { + Py_DECREF(maxsize_O); + return NULL; + } + typed = err; + } else + typed = 0; + + obj = (lru_cache_object *)type->tp_alloc(type, 0); + if (obj == NULL) { + Py_DECREF(maxsize_O); + return NULL; + } + + obj->root.prev = &obj->root; + obj->root.next = &obj->root; + obj->maxsize = maxsize; + obj->maxsize_O = maxsize_O; + if (!(obj->cache = PyDict_New())) { + Py_DECREF(obj); + Py_DECREF(maxsize_O); + return NULL; + } + obj->func = func; + Py_INCREF(func); + obj->wrapper = wrapper; + obj->misses = obj->hits = 0; + obj->typed = typed; + + return (PyObject *)obj; +} + +static void +lru_cache_clear_list(lru_list_elem *root) +{ + lru_list_elem *link = root->next; + while (link != root) { + lru_list_elem *next = link->next; + Py_DECREF(link->result); + PyMem_Free(link); + link = next; + } +} + +static void +lru_cache_dealloc(lru_cache_object *obj) +{ + Py_XDECREF(obj->maxsize_O); + Py_XDECREF(obj->func); + Py_XDECREF(obj->cache); + Py_XDECREF(obj->dict); + lru_cache_clear_list(&obj->root); + Py_TYPE(obj)->tp_free(obj); +} + +static PyObject * +lru_cache_call(lru_cache_object *self, PyObject *args, PyObject *kwds) +{ + return self->wrapper(self, args, kwds); +} + +static PyObject * +lru_cache_cache_info(lru_cache_object *self, PyObject *unused) +{ + PyObject *ret, *functools = PyImport_ImportModuleNoBlock("functools"); + if (functools == NULL) + return NULL; + ret = PyObject_CallMethod(functools, "_CacheInfo", "nnOn", + self->hits, self->misses, self->maxsize_O, + PyDict_Size(self->cache)); + Py_DECREF(functools); + return ret; +} + +static PyObject * +lru_cache_cache_clear(lru_cache_object *self, PyObject *unused) +{ + PyDict_Clear(self->cache); + self->hits = self->misses = 0; + lru_cache_clear_list(&self->root); + self->root.next = self->root.prev = &self->root; + Py_RETURN_NONE; +} + +static PyMethodDef lru_cache_methods[] = { + {"cache_info", (PyCFunction)lru_cache_cache_info, METH_NOARGS}, + {"cache_clear", (PyCFunction)lru_cache_cache_clear, METH_NOARGS}, + {NULL} +}; + +static PyGetSetDef lru_cache_getsetlist[] = { + {"__dict__", PyObject_GenericGetDict, PyObject_GenericSetDict}, + {NULL} +}; + +static PyTypeObject lru_cache_type = { + PyVarObject_HEAD_INIT(NULL, 0) + "functools._lru_cache", /* tp_name */ + sizeof(lru_cache_object), /* tp_basicsize */ + 0, /* tp_itemsize */ + /* methods */ + (destructor)lru_cache_dealloc, /* tp_dealloc */ + 0, /* tp_print */ + 0, /* tp_getattr */ + 0, /* tp_setattr */ + 0, /* tp_reserved */ + 0, /* tp_repr */ + 0, /* tp_as_number */ + 0, /* tp_as_sequence */ + 0, /* tp_as_mapping */ + 0, /* tp_hash */ + (ternaryfunc)lru_cache_call, /* tp_call */ + 0, /* tp_str */ + 0, /* tp_getattro */ + 0, /* tp_setattro */ + 0, /* tp_as_buffer */ + Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, + /* tp_flags */ + 0, /* tp_doc */ + 0, /* tp_traverse */ + 0, /* tp_clear */ + 0, /* tp_richcompare */ + 0, /* tp_weaklistoffset */ + 0, /* tp_iter */ + 0, /* tp_iternext */ + lru_cache_methods, /* tp_methods */ + 0, /* tp_members */ + lru_cache_getsetlist, /* tp_getset */ + 0, /* tp_base */ + 0, /* tp_dict */ + 0, /* tp_descr_get */ + 0, /* tp_descr_set */ + offsetof(lru_cache_object, dict), /* tp_dictoffset */ + 0, /* tp_init */ + 0, /* tp_alloc */ + lru_cache_new, /* tp_new */ +}; + /* module level code ********************************************************/ PyDoc_STRVAR(module_doc, @@ -552,6 +946,11 @@ {NULL, NULL} /* sentinel */ }; +static void +module_free(void *m) +{ + Py_DECREF(kwd_mark); +} static struct PyModuleDef _functoolsmodule = { PyModuleDef_HEAD_INIT, @@ -562,7 +961,7 @@ NULL, NULL, NULL, - NULL + module_free, }; PyMODINIT_FUNC @@ -573,6 +972,7 @@ char *name; PyTypeObject *typelist[] = { &partial_type, + &lru_cache_type, NULL }; @@ -580,6 +980,11 @@ if (m == NULL) return NULL; + if (!(kwd_mark = PyObject_CallObject((PyObject *)&PyBaseObject_Type, NULL))) { + Py_DECREF(m); + return NULL; + } + for (i=0 ; typelist[i] != NULL ; i++) { if (PyType_Ready(typelist[i]) < 0) { Py_DECREF(m);