From 6b779f29b05e3461f6bef8208df7bb9a79138214 Mon Sep 17 00:00:00 2001
From: Boris Besemer <borisbesemer@gmail.com>
Date: Thu, 22 Dec 2016 12:35:03 +0100
Subject: [PATCH] fixes a few issues

---
 src/personalisation/adapters.py      | 29 ++++++++++++++--------------
 src/personalisation/app_settings.py  |  5 ++---
 src/personalisation/utils.py         |  6 ------
 src/personalisation/wagtail_hooks.py |  2 +-
 4 files changed, 18 insertions(+), 24 deletions(-)

diff --git a/src/personalisation/adapters.py b/src/personalisation/adapters.py
index fa8259a..4979b5e 100644
--- a/src/personalisation/adapters.py
+++ b/src/personalisation/adapters.py
@@ -1,9 +1,14 @@
 import time
 
+from django.db.models import F
+
 from personalisation.models import AbstractBaseRule, Segment
 
 
 class BaseSegmentsAdapter(object):
+    def setup(self):
+        return None
+
     def _test_rules(self, rules, request):
         if len(rules) > 0:
             for rule in rules:
@@ -23,20 +28,18 @@ class SessionSegmentsAdapter(BaseSegmentsAdapter):
     def setup(self, request):
         self.request = request
 
-        # Set up segments dictionary object in the session
-        if 'segments' not in self.request.session:
-            self.request.session['segments'] = []
+        self.request.session.setdefault('segments', [])
 
-    def get(self):
+    def get_all_segments(self):
         return self.request.session['segments']
 
+    def get_segment(self, segment_id):
+        return next(item for item in self.request.session['segments'] if item.id == segment_id)
+
     def add(self, segment):
         def check_if_segmented(item):
             """Check if the user has been segmented"""
-            for seg in self.request.session['segments']:
-                if seg['encoded_name'] == item.encoded_name():
-                    return True
-            return False
+            return any(seg['encoded_name'] == item.encoded_name for seg in self.request.session['segments'])
 
         if not check_if_segmented(segment):
             segdict = {
@@ -56,15 +59,13 @@ class SessionSegmentsAdapter(BaseSegmentsAdapter):
 
         self.request.session['segments'] = current_segments
 
-        segments = Segment.objects.all().filter(status='enabled')
+        segments = Segment.objects.filter(status='enabled').prefetch_related('rules')
 
         for segment in segments:
             rules = AbstractBaseRule.__subclasses__()
             segment_rules = []
             for rule in rules:
-                queried_rules = rule.objects.filter(segment=segment)
-                for result in queried_rules:
-                    segment_rules.append(result)
+                segment_rules += rule.objects.filter(segment=segment)
             result = self._test_rules(segment_rules, self.request)
 
             if result:
@@ -73,11 +74,11 @@ class SessionSegmentsAdapter(BaseSegmentsAdapter):
 
         for seg in self.request.session['segments']:
             segment = Segment.objects.get(pk=seg['id'])
-            segment.visit_count = segment.visit_count + 1
+            segment.visit_count = F('visit_count') + 1
             segment.save()
 
     def check_segment_exists(self, segment):
         segments = self.request.session['segments']
 
-        return any(item for item in self.request.session['segments'] if segment.pk == item.id)
+        return any(item for item in segments if segment.pk == item.id)
 
diff --git a/src/personalisation/app_settings.py b/src/personalisation/app_settings.py
index 093b967..a64ceec 100644
--- a/src/personalisation/app_settings.py
+++ b/src/personalisation/app_settings.py
@@ -1,7 +1,6 @@
 from django.conf import settings
-
-from personalisation.utils import import_class
+from django.utils.module_loading import import_string
 
 
-segments_adapter = import_class(getattr(settings, 'PERSONALISATION_SEGMENTS_ADAPTER', 'personalisation.adapters.SessionSegmentsAdapter'))()
+segments_adapter = import_string(getattr(settings, 'PERSONALISATION_SEGMENTS_ADAPTER', 'personalisation.adapters.SessionSegmentsAdapter'))()
 
diff --git a/src/personalisation/utils.py b/src/personalisation/utils.py
index ab01ee2..c58ef30 100644
--- a/src/personalisation/utils.py
+++ b/src/personalisation/utils.py
@@ -1,11 +1,5 @@
-import importlib
-
 def impersonate_other_page(page, other_page):
     page.path = other_page.path
     page.depth = other_page.depth
     page.url_path = other_page.url_path
     page.title = other_page.title
-
-def import_class(name):
-    module_name, class_name = name.rsplit('.', 1)
-    return getattr(importlib.import_module(module_name), class_name)
diff --git a/src/personalisation/wagtail_hooks.py b/src/personalisation/wagtail_hooks.py
index 36a1e24..6428b60 100644
--- a/src/personalisation/wagtail_hooks.py
+++ b/src/personalisation/wagtail_hooks.py
@@ -88,7 +88,7 @@ def segment_user(page, request, serve_args, serve_kwargs):
 def serve_variation(page, request, serve_args, serve_kwargs):
     user_segments = []
 
-    for segment in segments_adapter.get():
+    for segment in segments_adapter.get_all_segments():
         try:
             user_segment = Segment.objects.get(pk=segment['id'],
                                                status='enabled')