Added RepresentativeViewMixin

A mixin to encapsulate pre-fetching logic for Representative country and
main_mandate.

representatives/contrib/parltrack/tests/test_import_representatives.py

@@ -12,7 +12,7 @@ def test_parltrack_import_representatives():
     fixture = os.path.join(os.path.dirname(__file__),
             'representatives_fixture.json')
     expected = os.path.join(os.path.dirname(__file__),
-            'representatives_expected.json')
+            '..', '..', '..', 'fixtures', 'representatives_test.json')
 
     # Disable django auto fields
     exclude = ('id', '_state', 'created', 'updated')

representatives/models.py

 # coding: utf-8
 
-# This file is part of compotista.
-#
-# compotista is free software: you can redistribute it and/or modify
-# it under the terms of the GNU Affero General Public License as
-# published by the Free Software Foundation, either version 3 of
-# the License, or any later version.
-#
-# compotista is distributed in the hope that it will
-# be useful, but WITHOUT ANY WARRANTY; without even the implied
-# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
-# See the GNU Affero General Public License for more details.
-#
-# You should have received a copy of the GNU General Affero Public
-# License along with django-representatives.
-# If not, see <http://www.gnu.org/licenses/>.
-#
-# Copyright (C) 2013 Laurent Peuch <cortex@worlddomination.be>
-# Copyright (C) 2015 Arnaud Fabre <af@laquadrature.net>
-
 import hashlib
-from datetime import datetime
+from datetime import datetime, date
 
 from django.db import models
 from django.utils.encoding import smart_str, smart_unicode

representatives/tests/test_views.py

+from django.test import TestCase
+
+from representatives.models import Representative
+from representatives.views import RepresentativeViewMixin
+
+
+class RepresentativeManagerTest(TestCase):
+    fixtures = ['representatives_test.json']
+
+    def test_prefetch_profile(self):
+        test = RepresentativeViewMixin()
+        reps = test.prefetch_for_representative_country_and_main_mandate(
+            Representative.objects.order_by('pk'))
+
+        with self.assertNumQueries(2):
+            # Cast to list to avoid [index] to cast a select with an offset
+            # below !
+            reps = [test.add_representative_country_and_main_mandate(r)
+                    for r in reps]
+
+            assert reps[0].country.code == 'AT'
+            assert reps[0].main_mandate is None
+
+            assert reps[1].country.code == 'SE'
+            assert reps[1].main_mandate.pk == 15

representatives/views.py

+import datetime
+
+from django.db import models
+
+from .models import Mandate
+
+
+class RepresentativeViewMixin(object):
+    """
+    A view mixin to add pre-fetched main_mandate and country to Representative
+
+    If a Representative was fetched from a QuerySet that have been through
+    prefetch_for_representative_country_and_main_mandate(), then
+    add_representative_country_and_main_mandate(representative) adds the
+    ``.country`` and ``.main_mandate`` properties "for free" - the prefetch
+    methods adds an extra query, but gets all.
+    """
+
+    def prefetch_for_representative_country_and_main_mandate(self, queryset):
+        """
+        Prefetch Mandates with their Group and Constituency with Country.
+        """
+        mandates = Mandate.objects.order_by(
+            '-end_date').select_related('constituency__country', 'group')
+        return queryset.prefetch_related(
+            models.Prefetch('mandates', queryset=mandates))
+
+    def add_representative_country_and_main_mandate(self, representative):
+        """
+        Set representative country and main_mandate.
+
+        Note that this will butcher your database if you don't use
+        self.prefetch_related.
+        """
+        today = datetime.date.today()
+
+        representative.country = None
+        representative.main_mandate = None
+
+        for m in representative.mandates.all():
+            if m.constituency.country_id and not representative.country:
+                representative.country = m.constituency.country
+
+            if (m.end_date > today and m.group.kind == 'group' and
+                not representative.main_mandate):
+
+                representative.main_mandate = m
+
+            if representative.country and not representative.main_mandate:
+                break
+
+        return representative