Add support for a 'country' filter in query parameters

mapit/templates/mapit/api/areas.html

@@ -23,6 +23,7 @@ <h3>{% blocktrans trimmed %}
                 <li><i>generation</i>, {% trans "to return areas in that generation (type and name lookups only)" %}.</li>
                 <li><i>min_generation</i>, {% trans "to return areas since that generation (type and name lookups only)" %}.</li>
                 <li><i>type</i>, {% trans "to restrict results to a type or types (multiple separated by commas; name lookup only)" %}.</li>
+                <li><i>country</i>, {% trans "to restrict results to areas with particular country codes (multiple separated by commas; type and name lookups only)" %}.</li>
             </ul></dd>
 
             <dt>{% trans "Returns" %}:</dt>

mapit/templates/mapit/api/point.html

@@ -58,6 +58,11 @@ <h3>{% blocktrans trimmed %}
                                 <i>min_generation</i>, to return results since that generation.
                             {% endblocktrans %}
                         </li>
+                        <li>
+                            {% blocktrans trimmed %}
+                                <i>country</i>, to restrict results to areas with particular country codes (multiple country codes separated by commas).
+                            {% endblocktrans %}
+                        </li>
                     </ul>
                 </dd>
 

mapit/tests/test_query_args.py

@@ -0,0 +1,126 @@
+# coding=utf-8
+
+from django.test import TestCase
+
+from mapit.models import Generation
+from mapit.views.areas import query_args
+
+
+class FakeRequest(object):
+    def __init__(self, params):
+        self.GET = params
+
+
+class QueryArgsTest(TestCase):
+
+    def setUp(self):
+        self.old_generation = Generation.objects.create(
+            active=False,
+            description="Old test generation",
+        )
+        self.active_generation = Generation.objects.create(
+            active=True,
+            description="Test generation",
+        )
+
+    def test_no_type(self):
+        args = query_args(FakeRequest({}), 'json', None)
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.active_generation.id,
+                'generation_low__lte': self.active_generation.id,
+            }
+        )
+
+    def test_one_type_in_query(self):
+        args = query_args(FakeRequest({'type': 'WMC'}), 'json', None)
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.active_generation.id,
+                'generation_low__lte': self.active_generation.id,
+                'type__code': 'WMC',
+            }
+        )
+
+    def test_two_types_in_query(self):
+        args = query_args(FakeRequest({'type': 'WMC,EUR'}), 'json', None)
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.active_generation.id,
+                'generation_low__lte': self.active_generation.id,
+                'type__code__in': ['WMC', 'EUR'],
+            }
+        )
+
+    def test_one_type_in_params(self):
+        args = query_args(FakeRequest({}), 'json', 'WMC')
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.active_generation.id,
+                'generation_low__lte': self.active_generation.id,
+                'type__code': 'WMC',
+            }
+        )
+
+    def test_two_types_in_params(self):
+        args = query_args(FakeRequest({}), 'json', 'WMC,EUR')
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.active_generation.id,
+                'generation_low__lte': self.active_generation.id,
+                'type__code__in': ['WMC', 'EUR'],
+            }
+        )
+
+    def test_generation_specified(self):
+        args = query_args(
+            FakeRequest({'generation': self.old_generation.id}),
+            'json',
+            None)
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.old_generation.id,
+                'generation_low__lte': self.old_generation.id,
+            }
+        )
+
+    def test_min_generation_specified(self):
+        args = query_args(
+            FakeRequest({'min_generation': self.old_generation.id}),
+            'json',
+            None)
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.old_generation.id,
+                'generation_low__lte': self.active_generation.id,
+            }
+        )
+
+    def test_one_country_in_query(self):
+        args = query_args(FakeRequest({'country': 'DE'}), 'json', None)
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.active_generation.id,
+                'generation_low__lte': self.active_generation.id,
+                'country__code': 'DE',
+            }
+        )
+
+    def test_two_countries_in_query(self):
+        args = query_args(FakeRequest({'country': 'DE,FR'}), 'json', None)
+        self.assertEqual(
+            args,
+            {
+                'generation_high__gte': self.active_generation.id,
+                'generation_low__lte': self.active_generation.id,
+                'country__code__in': ['DE', 'FR'],
+            }
+        )

mapit/views/areas.py

@@ -12,7 +12,7 @@
 from django.shortcuts import redirect, render
 from django.views.decorators.csrf import csrf_exempt
 
-from mapit.models import Area, Generation, Geometry, Code, Name
+from mapit.models import Area, Generation, Code, Name
 from mapit.shortcuts import output_json, output_html, output_polygon, get_object_or_404, set_timeout
 from mapit.middleware import ViewException
 from mapit.ratelimitcache import ratelimit
@@ -64,17 +64,22 @@ def query_args(request, format, type=None):
 
     if type is None:
         type = request.GET.get('type', '')
+    country_code = request.GET.get('country', '')
 
     args = {}
     if min_generation > -1:
         args = {
             'generation_low__lte': generation,
             'generation_high__gte': min_generation,
         }
-    if ',' in type:
-        args['type__code__in'] = type.split(',')
-    elif type:
-        args['type__code'] = type
+    for attr, value in [
+            ('type', type),
+            ('country', country_code),
+    ]:
+        if ',' in value:
+            args[attr + '__code__in'] = value.split(',')
+        elif value:
+            args[attr + '__code'] = value
 
     return args
 
@@ -358,29 +363,12 @@ def areas_by_point(request, srid, x, y, bb=False, format='json'):
     method = 'box' if bb and bb != 'polygon' else 'polygon'
 
     args = query_args(request, format)
-    type = request.GET.get('type', '')
-
-    if type and method == 'polygon':
-        args = dict(("area__%s" % k, v) for k, v in args.items())
-        # So this is odd. It doesn't matter if you specify types, PostGIS will
-        # do the contains test on all the geometries matching the bounding-box
-        # index, even if it could be much quicker to filter some out first
-        # (ie. the EUR ones).
-        args['polygon__bbcontains'] = location
-        shapes = Geometry.objects.filter(**args).defer('polygon')
-        areas = []
-        for shape in shapes:
-            try:
-                areas.append(Area.objects.get(polygons__id=shape.id, polygons__polygon__contains=location))
-            except:
-                pass
+
+    if method == 'box':
+        args['polygons__polygon__bbcontains'] = location
     else:
-        if method == 'box':
-            args['polygons__polygon__bbcontains'] = location
-        else:
-            geoms = list(Geometry.objects.filter(polygon__contains=location).defer('polygon'))
-            args['polygons__in'] = geoms
-        areas = Area.objects.filter(**args)
+        args['polygons__polygon__contains'] = location
+    areas = Area.objects.filter(**args)
 
     return output_areas(
         request, _('Areas covering the point ({0},{1})').format(x, y),