diff --git a/core/models.py b/core/models.py index c84f785..60ef492 100755 --- a/core/models.py +++ b/core/models.py @@ -1,13 +1,12 @@ import os import re -from itertools import chain from operator import itemgetter from bs4 import BeautifulSoup from django.core.exceptions import ObjectDoesNotExist from django.core.paginator import EmptyPage, PageNotAnInteger, Paginator from django.db import models -from django.db.models import Count, Q +from django.db.models import Case, Count, Q, Value, When from django.utils.encoding import python_2_unicode_compatible from django.utils.html import mark_safe from modelcluster.fields import ParentalKey @@ -344,37 +343,27 @@ def og_image(self): return image def children(self): - ordering = self.SITES_ORDERING[self.sites_ordering]['ordering'] + user_ordering = self.SITES_ORDERING[self.sites_ordering]['ordering'] + pages = WagtailSitePage.objects.live().filter(Q(path__startswith=self.path) | Q(in_cooperation_with=self)) # When ordering by `path`, the collaborations would either all be listed first or last # depending on whether the collaborator(s) page(s) was created before or after this page. # Adding an overwrite here so collaborations always appear last. if self.sites_ordering == self.SITES_ORDERING_PATH: - own_site_pages = WagtailSitePage.objects\ - .live()\ - .descendant_of(self)\ - .order_by(*ordering) - - collab_site_pages = WagtailSitePage.objects\ - .live()\ - .filter(in_cooperation_with=self)\ - .order_by(*ordering) - - # Using `itertools.chain` to chain both querysets - # as joining them with `|` or `union` does not produce the expected result. - site_pages = chain(own_site_pages, collab_site_pages) + pages = pages.annotate( + is_own=Case( + When(path__startswith=self.path, then=Value(True)), + default_value=Value(False), + output_field=models.BooleanField(), + ) + ).order_by('is_own', *user_ordering) # When ordering alphabetically or by creation date, # own sites and collaboration sites will be sorted together. else: - site_pages = WagtailSitePage.objects\ - .live()\ - .filter(Q(path__startswith=self.path) | Q(in_cooperation_with=self))\ - .order_by(*ordering) - - # At this point, `site_pages` might be an iterator or a QuerySet. - # It's not ideal to cast it to a list but at least it's consistent. - return list(site_pages) + pages = pages.order_by(*user_ordering) + + return pages def get_context(self, request, *args, **kwargs): # Get pages