Improve 'used by/usage count' implementation

- used_by() generates a database query instead of looping over a list of
  objects.
- usage_count() removed. Call count() on the queryset returned by
  used_by()
- Changed USAGE_COUNT to WAGTAIL_USAGE_COUNT_ENABLED. Added an inclusion
  tag for it.

wagtail/wagtailadmin/templates/wagtailadmin/shared/header.html

@@ -1,3 +1,4 @@
+{% load wagtailadmin_tags %}
 <header class="nice-padding {% if merged %}merged{% endif %} {% if tabbed %}tab-merged{% endif %} {% if search_form %}hasform{% endif %}">
     <div class="row">
         <div class="left">
@@ -16,9 +17,10 @@
             {% endif %}
         </div>
         <div class="right">
-            {% if usage_count or usage_count == 0 %}
+            {% usage_count_enabled as uc_enabled %}
+            {% if uc_enabled and usage_object %}
                 <div class="usagecount">
-                  <a href="{{ usage_url }}">Used {{ usage_count }} Time{% if not usage_count == 1 %}s{% endif %}</a>
+                  <a href="{{ usage_object.usage_url }}">Used {{ usage_object.used_by.count }} Time{% if not usage_count == 1 %}s{% endif %}</a>
                 </div>
             {% endif %}
             {% if add_link %}

wagtail/wagtailadmin/templatetags/wagtailadmin_tags.py

@@ -1,5 +1,6 @@
 from __future__ import unicode_literals
 
+from django.conf import settings
 from django import template
 from django.core import urlresolvers
 from django.utils.translation import ugettext_lazy as _
@@ -121,3 +122,13 @@ def hook_output(hook_name):
     """
     snippets = [fn() for fn in hooks.get_hooks(hook_name)]
     return ''.join(snippets)
+
+
+@register.assignment_tag
+def usage_count_enabled():
+    if hasattr(
+            settings, 'WAGTAIL_USAGE_COUNT_ENABLED'
+    ) or settings.WAGTAIL_USAGE_COUNT_ENABLED:
+        return True
+    else:
+        return False

wagtail/wagtailadmin/utils.py

@@ -1,5 +1,3 @@
-from django.conf import settings
-
 from modelcluster.fields import ParentalKey
 
 from wagtail.wagtailcore.models import Page
@@ -8,39 +6,31 @@ from wagtail.wagtailcore.models import Page
 def used_by(self):
     """Returns the pages that an object was used in."""
 
-    if not hasattr(settings, 'USAGE_COUNT') or not settings.USAGE_COUNT:
-        return []
+    pages = Page.objects.none()
 
-    related_objects = []
-    result = []
-
-    relations = self._meta.get_all_related_objects(
+    # get all the relation objects for self
+    relations = type(self)._meta.get_all_related_objects(
         include_hidden=True,
         include_proxy_eq=True
     )
     for relation in relations:
-        related_objects.extend(list(relation.model._base_manager.filter(
-            **{relation.field.name: self.id}
-        )))
-    for r in related_objects:
-        if isinstance(r, Page):
-            result.append(r)
+        # if the relation is between self and a page, get the page
+        if issubclass(relation.model, Page):
+            pages |= Page.objects.filter(
+                id__in=relation.model._base_manager.filter(**{
+                    relation.field.name: self.id
+                }).values_list('id', flat=True)
+            )
         else:
-            parental_keys = get_parental_keys(r)
-            for parental_key in parental_keys:
-                result.append(getattr(r, parental_key.name))
+        # if the relation is between self and an object that has a page as a
+        # property, return the page
+            for f in relation.model._meta.fields:
+                if isinstance(f, ParentalKey) and issubclass(f.rel.to, Page):
+                    pages |= Page.objects.filter(
+                        id__in=relation.model._base_manager.filter(
+                            **{
+                                relation.field.name: self.id
+                            }).values_list(f.attname, flat=True)
+                    )
 
-    return result
-
-
-def usage_count(self):
-    """Returns the number of times that an obect has been used in a page"""
-
-    if not hasattr(settings, 'USAGE_COUNT') or not settings.USAGE_COUNT:
-        return None
-
-    return len(used_by(self))
-
-
-def get_parental_keys(obj):
-    return [field for field in obj._meta.fields if type(field) == ParentalKey]
+    return pages

wagtail/wagtaildocs/models.py

@@ -12,7 +12,7 @@ from django.utils.translation import ugettext_lazy  as _
 from django.utils.encoding import python_2_unicode_compatible
 
 from wagtail.wagtailadmin.taggable import TagSearchable
-from wagtail.wagtailadmin.utils import usage_count, used_by
+from wagtail.wagtailadmin.utils import used_by
 from wagtail.wagtailsearch import indexed
 
 
@@ -48,14 +48,15 @@ class Document(models.Model, TagSearchable):
     def url(self):
         return reverse('wagtaildocs_serve', args=[self.id, self.filename])
 
-    @property
-    def usage_count(self):
-        return usage_count(self)
-
     @property
     def used_by(self):
         return used_by(self)
 
+    @property
+    def usage_url(self):
+        return reverse('wagtaildocs_document_usage',
+                       args=(self.id,))
+
     def is_editable_by_user(self, user):
         if user.has_perm('wagtaildocs.change_document'):
             # user has global permission to change documents

wagtail/wagtaildocs/templates/wagtaildocs/documents/edit.html

@@ -13,7 +13,7 @@
 
 {% block content %}
     {% trans "Editing" as editing_str %}
-    {% include "wagtailadmin/shared/header.html" with title=editing_str subtitle=document.title icon="doc-full-inverse" usage_count=document.usage_count usage_url=usage_url %}
+    {% include "wagtailadmin/shared/header.html" with title=editing_str subtitle=document.title icon="doc-full-inverse" usage_object=document %}
 
     <div class="nice-padding">
         <form action="{% url 'wagtaildocs_edit_document' document.id %}" method="POST" enctype="multipart/form-data">

wagtail/wagtaildocs/views/documents.py

@@ -102,9 +102,6 @@ def edit(request, document_id):
     if not doc.is_editable_by_user(request.user):
         raise PermissionDenied
 
-    usage_url = reverse('wagtaildocs_document_usage',
-                        args=(doc.id,))
-
     if request.POST:
         original_file = doc.file
         form = DocumentForm(request.POST, request.FILES, instance=doc)
@@ -124,8 +121,7 @@ def edit(request, document_id):
 
     return render(request, "wagtaildocs/documents/edit.html", {
         'document': doc,
-        'form': form,
-        'usage_url': usage_url
+        'form': form
     })
 
 

wagtail/wagtailimages/models.py

@@ -15,6 +15,7 @@ from django.utils.html import escape, format_html_join
 from django.conf import settings
 from django.utils.translation import ugettext_lazy  as _
 from django.utils.encoding import python_2_unicode_compatible
+from django.core.urlresolvers import reverse
 
 from unidecode import unidecode
 
@@ -22,7 +23,7 @@ from wagtail.wagtailadmin.taggable import TagSearchable
 from wagtail.wagtailimages.backends import get_image_backend
 from wagtail.wagtailsearch import indexed
 from .utils import validate_image_format
-from wagtail.wagtailadmin.utils import usage_count, used_by
+from wagtail.wagtailadmin.utils import used_by
 
 
 @python_2_unicode_compatible
@@ -50,14 +51,15 @@ class AbstractImage(models.Model, TagSearchable):
 
     tags = TaggableManager(help_text=None, blank=True, verbose_name=_('Tags'))
 
-    @property
-    def usage_count(self):
-        return usage_count(self)
-
     @property
     def used_by(self):
         return used_by(self)
 
+    @property
+    def usage_url(self):
+        return reverse('wagtailimages_image_usage',
+                       args=(self.id,))
+
     search_fields = TagSearchable.search_fields + (
         indexed.FilterField('uploaded_by_user'),
     )

wagtail/wagtailimages/templates/wagtailimages/images/edit.html

@@ -13,7 +13,7 @@
 
 {% block content %}
     {% trans "Editing" as editing_str %}
-    {% include "wagtailadmin/shared/header.html" with title=editing_str subtitle=image.title icon="image" usage_count=image.usage_count usage_url=usage_url %}
+    {% include "wagtailadmin/shared/header.html" with title=editing_str subtitle=image.title icon="image" usage_object=image %}
 
     <div class="row row-flush nice-padding">
 

wagtail/wagtailimages/views/images.py

@@ -5,7 +5,6 @@ from django.contrib.auth.decorators import permission_required
 from django.core.exceptions import PermissionDenied
 from django.utils.translation import ugettext as _
 from django.views.decorators.vary import vary_on_headers
-from django.core.urlresolvers import reverse
 
 from wagtail.wagtailadmin.forms import SearchForm
 
@@ -33,7 +32,6 @@ def index(request):
         if form.is_valid():
             query_string = form.cleaned_data['q']
 
-            is_searching = True
             if not request.user.has_perm('wagtailimages.change_image'):
                 # restrict to the user's own images
                 images = Image.search(query_string, filters={'uploaded_by_user_id': request.user.id})
@@ -81,9 +79,6 @@ def edit(request, image_id):
     if not image.is_editable_by_user(request.user):
         raise PermissionDenied
 
-    usage_url = reverse('wagtailimages_image_usage',
-                        args=(image.id,))
-
     if request.POST:
         original_file = image.file
         form = ImageForm(request.POST, request.FILES, instance=image)
@@ -104,8 +99,7 @@ def edit(request, image_id):
 
     return render(request, "wagtailimages/images/edit.html", {
         'image': image,
-        'form': form,
-        'usage_url': usage_url
+        'form': form
     })
 
 

wagtail/wagtailsnippets/models.py

@@ -1,5 +1,7 @@
 from django.contrib.contenttypes.models import ContentType
-from wagtail.wagtailadmin.utils import usage_count, used_by
+from django.core.urlresolvers import reverse
+
+from wagtail.wagtailadmin.utils import used_by
 
 SNIPPET_MODELS = []
 
@@ -19,7 +21,15 @@ def get_snippet_content_types():
 
 def register_snippet(model):
     if model not in SNIPPET_MODELS:
-        model.usage_count = usage_count
         model.used_by = used_by
+        model.usage_url = get_snippet_usage_url
         SNIPPET_MODELS.append(model)
         SNIPPET_MODELS.sort(key=lambda x: x._meta.verbose_name)
+
+
+def get_snippet_usage_url(self):
+    content_type = ContentType.objects.get_for_model(self)
+    return reverse('wagtailsnippets_usage',
+                   args=(content_type.app_label,
+                         content_type.model,
+                         self.id,))

wagtail/wagtailsnippets/templates/wagtailsnippets/snippets/edit.html

@@ -4,7 +4,7 @@
 {% block bodyclass %}menu-snippets{% endblock %}
 {% block content %}
     {% trans "Editing" as editing_str %}
-    {% include "wagtailadmin/shared/header.html" with title=editing_str subtitle=instance icon="snippet" usage_count=instance.usage_count usage_url=usage_url %}
+    {% include "wagtailadmin/shared/header.html" with title=editing_str subtitle=instance icon="snippet" usage_object=instance %}
 
     <form action="{% url 'wagtailsnippets_edit' content_type.app_label content_type.model instance.id %}" method="POST">
         {% csrf_token %}

wagtail/wagtailsnippets/views/snippets.py

@@ -158,11 +158,6 @@ def edit(request, content_type_app_name, content_type_model_name, id):
     edit_handler_class = get_snippet_edit_handler(model)
     form_class = edit_handler_class.get_form_class(model)
 
-    usage_url = reverse('wagtailsnippets_usage',
-                        args=(content_type_app_name,
-                              content_type_model_name,
-                              id))
-
     if request.POST:
         form = form_class(request.POST, request.FILES, instance=instance)
 
@@ -188,8 +183,7 @@ def edit(request, content_type_app_name, content_type_model_name, id):
         'content_type': content_type,
         'snippet_type_name': snippet_type_name,
         'instance': instance,
-        'edit_handler': edit_handler,
-        'usage_url': usage_url
+        'edit_handler': edit_handler
     })