move complicated Location API Queryset into LocationSlug

src/c3nav/mapdata/api.py

@@ -1,9 +1,9 @@
 import mimetypes
 import operator
-from functools import reduce, wraps
+from functools import wraps
 
 from django.core.cache import cache
-from django.db.models import Prefetch, Q
+from django.db.models import Prefetch
 from django.http import HttpResponse
 from django.shortcuts import redirect
 from django.utils.cache import get_conditional_response
@@ -223,29 +223,8 @@ class LocationViewSet(RetrieveModelMixin, GenericViewSet):
     queryset = LocationSlug.objects.all()
     lookup_field = 'slug'
 
-    def get_queryset(self, mode=None):
-        queryset = super().get_queryset().order_by('id')
-
-        conditions = []
-        for model in get_submodels(Location):
-            related_name = model._meta.default_related_name
-            condition = Q(**{related_name+'__isnull': False})
-            if mode == 'search':
-                condition &= Q(**{related_name+'__can_search': True})
-            elif mode == 'search-describe':
-                condition &= Q(**{related_name+'__can_search': True}) | Q(**{related_name+'__can_describe': True})
-            # noinspection PyUnresolvedReferences
-            condition &= model.q_for_request(self.request, prefix=related_name+'__')
-            conditions.append(condition)
-        queryset = queryset.filter(reduce(operator.or_, conditions))
-
-        # prefetch locationgroups
-        base_qs = LocationGroup.qs_for_request(self.request).select_related('category')
-        for model in get_submodels(SpecificLocation):
-            queryset = queryset.prefetch_related(Prefetch(model._meta.default_related_name + '__groups',
-                                                          queryset=base_qs))
-
-        return queryset
+    def get_queryset(self, can=None):
+        return LocationSlug.location_qs_for_request(self.request, can=can)
 
     @simple_api_cache()
     def list(self, request, *args, **kwargs):
@@ -265,7 +244,7 @@ class LocationViewSet(RetrieveModelMixin, GenericViewSet):
             )
             queryset = cache.get(queryset_cache_key, None)
             if queryset is None or 1:
-                queryset = self.get_queryset(mode=('searchable' if searchable else 'searchable-describe'))
+                queryset = self.get_queryset(can=(('search', ) if searchable else ('search', 'describe')))
 
                 queryset = tuple(obj.get_child() for obj in queryset)
 
@@ -316,7 +295,7 @@ class LocationViewSet(RetrieveModelMixin, GenericViewSet):
 
     @simple_api_cache()
     def retrieve(self, request, slug=None, *args, **kwargs):
-        result = Location.get_by_slug(slug, self.get_queryset())
+        result = Location.get_by_slug(slug, request)
         if result is None:
             raise NotFound
         result = result.get_child()

src/c3nav/mapdata/models/locations.py

@@ -1,9 +1,11 @@
+import operator
 from collections import OrderedDict
 from contextlib import suppress
+from functools import reduce
 
 from django.apps import apps
 from django.db import models
-from django.db.models import Prefetch
+from django.db.models import Prefetch, Q
 from django.utils.functional import cached_property
 from django.utils.text import format_lazy
 from django.utils.translation import ugettext_lazy as _
@@ -67,6 +69,29 @@ class LocationSlug(SerializableMixin, models.Model):
         verbose_name_plural = _('Location with Slug')
         default_related_name = 'locationslugs'
 
+    @classmethod
+    def location_qs_for_request(cls, request, can=None):
+        queryset = cls.objects.all().order_by('id')
+
+        conditions = []
+        for model in get_submodels(Location):
+            related_name = model._meta.default_related_name
+            condition = Q(**{related_name + '__isnull': False})
+            if can:
+                condition &= reduce(operator.or_, (Q(**{related_name+'__can_'+s: True}) for s in can))
+            # noinspection PyUnresolvedReferences
+            condition &= model.q_for_request(request, prefix=related_name + '__')
+            conditions.append(condition)
+        queryset = queryset.filter(reduce(operator.or_, conditions))
+
+        # prefetch locationgroups
+        base_qs = LocationGroup.qs_for_request(request).select_related('category')
+        for model in get_submodels(SpecificLocation):
+            queryset = queryset.prefetch_related(Prefetch(model._meta.default_related_name + '__groups',
+                                                          queryset=base_qs))
+
+        return queryset
+
 
 class Location(LocationSlug, AccessRestrictionMixin, TitledMixin, models.Model):
     can_search = models.BooleanField(default=True, verbose_name=_('can be searched'))
@@ -103,9 +128,11 @@ class Location(LocationSlug, AccessRestrictionMixin, TitledMixin, models.Model):
         return self.slug
 
     @classmethod
-    def get_by_slug(cls, slug, queryset=None):
-        if queryset is None:
+    def get_by_slug(cls, slug, request=None, can=None):
+        if request is None:
             queryset = LocationSlug.objects.all()
+        else:
+            queryset = LocationSlug.location_qs_for_request(request, can)
 
         if ':' in slug:
             code, pk = slug.split(':', 1)