diff --git a/apps/establishment/views/web.py b/apps/establishment/views/web.py index 1a0d5f58..0699d9d0 100644 --- a/apps/establishment/views/web.py +++ b/apps/establishment/views/web.py @@ -19,9 +19,12 @@ class EstablishmentMixinView: def get_queryset(self): """Overridden method 'get_queryset'.""" - return models.Establishment.objects.published() \ - .with_base_related() \ - .annotate_in_favorites(self.request.user) + qs = models.Establishment.objects.published() \ + .with_base_related() \ + .annotate_in_favorites(self.request.user) + if self.request.country_code: + qs = qs.by_country_code(self.request.country_code) + return qs class EstablishmentListView(EstablishmentMixinView, generics.ListAPIView): @@ -30,13 +33,6 @@ class EstablishmentListView(EstablishmentMixinView, generics.ListAPIView): filter_class = filters.EstablishmentFilter serializer_class = serializers.EstablishmentBaseSerializer - def get_queryset(self): - """Overridden method 'get_queryset'.""" - qs = super(EstablishmentListView, self).get_queryset() - if self.request.country_code: - qs = qs.by_country_code(self.request.country_code) - return qs - class EstablishmentRetrieveView(EstablishmentMixinView, generics.RetrieveAPIView): """Resource for getting a establishment.""" @@ -45,9 +41,7 @@ class EstablishmentRetrieveView(EstablishmentMixinView, generics.RetrieveAPIView serializer_class = serializers.EstablishmentDetailSerializer def get_queryset(self): - return super().get_queryset() \ - .by_country_code(self.request.country_code) \ - .with_extended_related() + return super().get_queryset().with_extended_related() class EstablishmentRecentReviewListView(EstablishmentListView):