diff --git a/representatives/contrib/parltrack/tests/test_import_representatives.py b/representatives/contrib/parltrack/tests/test_import_representatives.py index 5f443a48c5d8750a0e261d7a614a6eada7938f18..c05d7f69cb0d382f7c753bbdd93183fece55348e 100644 --- a/representatives/contrib/parltrack/tests/test_import_representatives.py +++ b/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') diff --git a/representatives/contrib/parltrack/tests/representatives_expected.json b/representatives/fixtures/representatives_test.json similarity index 100% rename from representatives/contrib/parltrack/tests/representatives_expected.json rename to representatives/fixtures/representatives_test.json diff --git a/representatives/models.py b/representatives/models.py index ad02f1475073eb1fc332659494aaa1164ccd215d..39db51cc9d067939064dabe63d0bf1b723dbefef 100644 --- a/representatives/models.py +++ b/representatives/models.py @@ -1,26 +1,7 @@ # 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 diff --git a/representatives/tests/test_views.py b/representatives/tests/test_views.py new file mode 100644 index 0000000000000000000000000000000000000000..95923fafc493610064c57679d4f38bfa37a943c5 --- /dev/null +++ b/representatives/tests/test_views.py @@ -0,0 +1,25 @@ +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 diff --git a/representatives/views.py b/representatives/views.py new file mode 100644 index 0000000000000000000000000000000000000000..0540fc94f062036ebadb8c674f36855e337c2407 --- /dev/null +++ b/representatives/views.py @@ -0,0 +1,52 @@ +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