From 417cafe69bd2ae19c85790fcbc938d190e701165 Mon Sep 17 00:00:00 2001
From: Tim Heap <tim@timheap.me>
Date: Fri, 30 Oct 2015 15:06:24 +1100
Subject: [PATCH] Add @cached_classmethod, use instead of
 @classmethod/@lru_cache

---
 wagtail/utils/decorators.py                   | 59 +++++++++++++++++++
 wagtail/wagtailadmin/edit_handlers.py         | 26 ++++----
 .../wagtailadmin/tests/test_edit_handlers.py  |  8 ++-
 3 files changed, 81 insertions(+), 12 deletions(-)
 create mode 100644 wagtail/utils/decorators.py

diff --git a/wagtail/utils/decorators.py b/wagtail/utils/decorators.py
new file mode 100644
index 0000000000..c640a030ed
--- /dev/null
+++ b/wagtail/utils/decorators.py
@@ -0,0 +1,59 @@
+from __future__ import absolute_import, unicode_literals
+
+from django.utils.functional import cached_property
+import functools
+
+
+# Need to inherit from object explicitly, to turn ``cached_classmethod`` in to
+# a new-style class. WeakKeyDictionary is an old-style class, which do not
+# support descriptors.
+class cached_classmethod(dict):
+    """
+    Cache the result of a no-arg class method.
+    .. code-block:: python
+        class Foo(object):
+            @cached_classmethod
+            def bar(cls):
+                # Some expensive computation
+                return 'baz'
+    Similar to ``@lru_cache``, but the cache is per-class, stores a single
+    value, and thus doesn't fill up; where as ``@lru_cache`` is global across
+    all classes, and could fill up if too many classes were used.
+    """
+
+    def __init__(self, fn):
+        self.fn = fn
+        functools.update_wrapper(self, fn)
+
+    def __get__(self, instance, owner):
+        """ Get the class_cache for this type when accessed """
+        return self[owner]
+
+    def __missing__(self, cls):
+        """ Make a new class_cache on cache misses """
+        value = _cache(self, cls, self.fn)
+        self[cls] = value
+        return value
+
+
+class _cache(object):
+    """ Calls the real class method behind when called, caching the result """
+    def __init__(self, cache, cls, fn):
+        self.cache = cache
+        self.cls = cls
+        self.fn = fn
+        functools.update_wrapper(self, fn)
+
+    @cached_property
+    def value(self):
+        """ Generate the cached value """
+        return self.fn(self.cls)
+
+    def __call__(self):
+        """ Get the cached value """
+        return self.value
+
+    def cache_clear(self):
+        """ Clear the cached value. """
+        # Named after lru_cache.cache_clear
+        self.cache.pop(self.cls, None)
diff --git a/wagtail/wagtailadmin/edit_handlers.py b/wagtail/wagtailadmin/edit_handlers.py
index 0bc8340018..e4c8a2a320 100644
--- a/wagtail/wagtailadmin/edit_handlers.py
+++ b/wagtail/wagtailadmin/edit_handlers.py
@@ -6,7 +6,6 @@ from django.contrib.contenttypes.models import ContentType
 from django.core.exceptions import ImproperlyConfigured
 from django.forms.models import fields_for_model
 from django.template.loader import render_to_string
-from django.utils.lru_cache import lru_cache
 from django.utils.safestring import mark_safe
 from django.utils.six import text_type
 from django.utils.translation import ugettext_lazy
@@ -16,6 +15,8 @@ from wagtail.wagtailcore.models import Page
 from wagtail.wagtailcore.utils import (
     camelcase_to_underscore, resolve_model_string)
 
+from wagtail.utils.decorators import cached_classmethod
+
 # DIRECT_FORM_FIELD_OVERRIDES, FORM_FIELD_OVERRIDES are imported for backwards
 # compatibility, as people are likely importing them from here and then
 # appending their own overrides
@@ -271,19 +272,22 @@ class BaseFormEditHandler(BaseCompositeEditHandler):
     # WagtailAdminModelForm
     base_form_class = WagtailAdminModelForm
 
+    _form_class = None
+
     @classmethod
-    @lru_cache()
     def get_form_class(cls, model):
         """
         Construct a form class that has all the fields and formsets named in
         the children of this edit handler.
         """
-        return get_form_for_model(
-            model,
-            form_class=cls.base_form_class,
-            fields=cls.required_fields(),
-            formsets=cls.required_formsets(),
-            widgets=cls.widget_overrides())
+        if cls._form_class is None:
+            cls._form_class = get_form_for_model(
+                model,
+                form_class=cls.base_form_class,
+                fields=cls.required_fields(),
+                formsets=cls.required_formsets(),
+                widgets=cls.widget_overrides())
+        return cls._form_class
 
 
 class BaseTabbedInterface(BaseFormEditHandler):
@@ -710,9 +714,11 @@ Page.settings_panels = [
 Page.base_form_class = WagtailAdminPageForm
 
 
-@classmethod
-@lru_cache()
+@cached_classmethod
 def get_edit_handler(cls):
+    """
+    Get the EditHandler to use in the Wagtail admin when editing this page type.
+    """
     if hasattr(cls, 'edit_handler'):
         return cls.edit_handler.bind_to_model(cls)
 
diff --git a/wagtail/wagtailadmin/tests/test_edit_handlers.py b/wagtail/wagtailadmin/tests/test_edit_handlers.py
index 6ee9fb943f..a62c5e4e01 100644
--- a/wagtail/wagtailadmin/tests/test_edit_handlers.py
+++ b/wagtail/wagtailadmin/tests/test_edit_handlers.py
@@ -429,8 +429,12 @@ class TestPageChooserPanel(TestCase):
 
     def test_render_js_init_with_can_choose_root_true(self):
         # construct an alternative page chooser panel object, with can_choose_root=True
-        MyPageChooserPanel = PageChooserPanel('page', can_choose_root=True).bind_to_model(PageChooserModel)
-        PageChooserForm = MyPageChooserPanel.get_form_class(PageChooserModel)
+
+        MyPageObjectList = ObjectList([
+            PageChooserPanel('page', can_choose_root=True)
+        ]).bind_to_model(PageChooserModel)
+        MyPageChooserPanel = MyPageObjectList.children[0]
+        PageChooserForm = MyPageObjectList.get_form_class(EventPageChooserModel)
 
         form = PageChooserForm(instance=self.test_instance)
         page_chooser_panel = MyPageChooserPanel(instance=self.test_instance, form=form)