From bb34bddaf47dc5ac2a99ca6847012afa46979d1e Mon Sep 17 00:00:00 2001
From: Paul J Stevens <p.stevens@labdigital.nl>
Date: Sat, 26 May 2018 12:01:26 +0200
Subject: [PATCH] add custom model manager

---
 src/wagtail_personalisation/models.py | 14 ++++++++++++++
 src/wagtail_personalisation/utils.py  | 15 +++++----------
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/src/wagtail_personalisation/models.py b/src/wagtail_personalisation/models.py
index 0e61e77..d88a1fd 100644
--- a/src/wagtail_personalisation/models.py
+++ b/src/wagtail_personalisation/models.py
@@ -1,10 +1,12 @@
 from __future__ import absolute_import, unicode_literals
+
 import random
 
 from django import forms
 from django.conf import settings
 from django.core.validators import MaxValueValidator, MinValueValidator
 from django.db import models, transaction
+from django.db.models import F
 from django.template.defaultfilters import slugify
 from django.utils.encoding import python_2_unicode_compatible
 from django.utils.functional import cached_property
@@ -199,6 +201,16 @@ class Segment(ClusterableModel):
         return False
 
 
+class PersonalisablePageManager(models.Manager):
+
+    def canonicals(self):
+        return (
+            self.filter(
+                personalisable_canonical_metadata__canonical_page_id=F(
+                    'personalisable_canonical_metadata__variant__id'))
+        )
+
+
 class PersonalisablePageMetadata(ClusterableModel):
     """The personalisable page model. Allows creation of variants with linked
     segments.
@@ -219,6 +231,8 @@ class PersonalisablePageMetadata(ClusterableModel):
         on_delete=models.SET_NULL,
         null=True, blank=True)
 
+    objects = PersonalisablePageManager()
+
     @cached_property
     def has_variants(self):
         """Return a boolean indicating whether or not the personalisable page
diff --git a/src/wagtail_personalisation/utils.py b/src/wagtail_personalisation/utils.py
index 6646027..59a1246 100644
--- a/src/wagtail_personalisation/utils.py
+++ b/src/wagtail_personalisation/utils.py
@@ -1,6 +1,5 @@
 import time
 
-from django.db.models import F
 from django.template.base import FilterExpression, kwarg_re
 from django.utils import timezone
 
@@ -99,13 +98,9 @@ def parse_tag(token, parser):
 def exclude_variants(pages):
     """Checks if page is not a variant
 
-    :param pages: List of pages to check
-    :type pages: list
-    :return: List of pages that aren't variants
-    :rtype: list
+    :param pages: Set of pages to check
+    :type pages: QuerySet
+    :return: Queryset of pages that aren't variants
+    :rtype: QuerySet
     """
-    return (
-        pages.filter(
-            personalisable_canonical_metadata__canonical_page_id=F(
-                'personalisable_canonical_metadata__variant__id'))
-    )
+    return pages.canonicals()