refactored methods for getting similar establishments

apps/establishment/models.py

@@ -221,15 +221,16 @@ class EstablishmentQuerySet(models.QuerySet):
         Return filtered QuerySet by base filters.
         Filters including:
         1 Filter by type (and subtype) establishment.
-        2 Filter by published Review.
-        3 With annotated distance.
+        2 With annotated distance.
+        3 By country
         """
         filters = {
-            'reviews__status': Review.READY,
             'establishment_type': establishment.establishment_type,
+            'address__city__country': establishment.address.city.country
         }
         if establishment.establishment_subtypes.exists():
             filters.update({'establishment_subtypes__in': establishment.establishment_subtypes.all()})
+
         return self.exclude(id=establishment.id) \
                    .filter(**filters) \
                    .annotate_distance(point=establishment.location)
@@ -248,28 +249,25 @@ class EstablishmentQuerySet(models.QuerySet):
                 .values('id')
         )
 
-    def similar_restaurants(self, slug):
+    def similar_restaurants(self, restaurant):
         """
         Return QuerySet with objects that similar to Restaurant.
-        :param slug: str restaurant slug
+        :param restaurant: Establishment instance.
         """
-        restaurant_qs = self.filter(slug=slug)
-        if restaurant_qs.exists():
-            restaurant = restaurant_qs.first()
-            ids_by_subquery = self.similar_base_subquery(
-                establishment=restaurant,
-                filters={
-                    'public_mark__gte': 10,
-                    'establishment_gallery__is_main': True,
-                }
-            )
-            return self.filter(id__in=ids_by_subquery) \
-                .annotate_intermediate_public_mark() \
-                .annotate_mark_similarity(mark=restaurant.public_mark) \
-                .order_by('mark_similarity') \
-                .distinct('mark_similarity', 'id')
-        else:
-            return self.none()
+
+        ids_by_subquery = self.similar_base_subquery(
+            establishment=restaurant,
+            filters={
+                'reviews__status': Review.READY,
+                'public_mark__gte': 10,
+                'establishment_gallery__is_main': True,
+            }
+        )
+        return self.filter(id__in=ids_by_subquery) \
+                   .annotate_intermediate_public_mark() \
+                   .annotate_mark_similarity(mark=restaurant.public_mark) \
+                   .order_by('mark_similarity') \
+                   .distinct('mark_similarity', 'id')
 
     def same_subtype(self, establishment):
         """Annotate flag same subtype."""
@@ -282,21 +280,17 @@ class EstablishmentQuerySet(models.QuerySet):
             output_field=models.BooleanField(default=False)
         ))
 
-    def similar_artisans_producers(self, slug):
+    def similar_artisans_producers(self, establishment):
         """
         Return QuerySet with objects that similar to Artisan/Producer(s).
-        :param slug: str artisan/producer slug
+        :param establishment: Establishment instance
         """
-        establishment_qs = self.filter(slug=slug)
-        if establishment_qs.exists():
-            establishment = establishment_qs.first()
-            return self.similar_base(establishment) \
-                       .same_subtype(establishment) \
-                       .order_by(F('same_subtype').desc(),
-                                 F('distance').asc()) \
-                       .distinct('same_subtype', 'distance', 'id')
-        else:
-            return self.none()
+        return self.similar_base(establishment) \
+                   .same_subtype(establishment) \
+                   .has_published_reviews() \
+                   .order_by(F('same_subtype').desc(),
+                             F('distance').asc()) \
+                   .distinct('same_subtype', 'distance', 'id')
 
     def by_wine_region(self, wine_region):
         """
@@ -312,23 +306,19 @@ class EstablishmentQuerySet(models.QuerySet):
         """
         return self.filter(wine_origin__wine_sub_region=wine_sub_region).distinct()
 
-    def similar_wineries(self, slug: str):
+    def similar_wineries(self, winery):
         """
         Return QuerySet with objects that similar to Winery.
-        :param establishment_slug: str Establishment slug
+        :param winery: Establishment instance
         """
-        winery_qs = self.filter(slug=slug)
-        if winery_qs.exists():
-            winery = winery_qs.first()
-            return self.similar_base(winery) \
-                       .order_by(F('wine_origins__wine_region').asc(),
-                                 F('wine_origins__wine_sub_region').asc()) \
-                       .annotate_distance(point=winery.location) \
-                       .order_by('distance') \
-                       .distinct('distance', 'wine_origins__wine_region',
-                                 'wine_origins__wine_sub_region', 'id')
-        else:
-            return self.none()
+        return self.similar_base(winery) \
+                   .order_by(F('wine_origins__wine_region').asc(),
+                             F('wine_origins__wine_sub_region').asc(),
+                             F('distance').asc()) \
+                   .distinct('wine_origins__wine_region',
+                             'wine_origins__wine_sub_region',
+                             'distance',
+                             'id')
 
     def last_reviewed(self, point: Point):
         """

apps/establishment/views/web.py

@@ -46,6 +46,19 @@ class EstablishmentSimilarView(EstablishmentListView):
     serializer_class = serializers.EstablishmentSimilarSerializer
     pagination_class = PortionPagination
 
+    def get_base_object(self):
+        """
+        Return base establishment instance for a getting list of similar establishments.
+        """
+        establishment = get_object_or_404(models.Establishment.objects.all(),
+                                          slug=self.kwargs.get('slug'))
+        return establishment
+
+    def get_queryset(self):
+        """Overridden get_queryset method."""
+        return EstablishmentMixinView.get_queryset(self) \
+                                     .has_location()
+
 
 class EstablishmentRetrieveView(EstablishmentMixinView, generics.RetrieveAPIView):
     """Resource for getting a establishment."""
@@ -88,9 +101,14 @@ class RestaurantSimilarListView(EstablishmentSimilarView):
 
     def get_queryset(self):
         """Overridden get_queryset method"""
-        return EstablishmentMixinView.get_queryset(self) \
-                                     .has_location() \
-                                     .similar_restaurants(slug=self.kwargs.get('slug'))
+        qs = super(RestaurantSimilarListView, self).get_queryset()
+        base_establishment = self.get_base_object()
+
+        if base_establishment:
+            return qs.similar_restaurants(base_establishment)
+        else:
+            return EstablishmentMixinView.get_queryset(self) \
+                                         .none()
 
 
 class WinerySimilarListView(EstablishmentSimilarView):
@@ -98,9 +116,13 @@ class WinerySimilarListView(EstablishmentSimilarView):
 
     def get_queryset(self):
         """Overridden get_queryset method"""
-        return EstablishmentMixinView.get_queryset(self) \
-                                     .has_location() \
-                                     .similar_wineries(slug=self.kwargs.get('slug'))
+        qs = EstablishmentSimilarView.get_queryset(self)
+        base_establishment = self.get_base_object()
+
+        if base_establishment:
+            return qs.similar_wineries(base_establishment)
+        else:
+            return qs.none()
 
 
 class ArtisanProducerSimilarListView(EstablishmentSimilarView):
@@ -108,9 +130,13 @@ class ArtisanProducerSimilarListView(EstablishmentSimilarView):
 
     def get_queryset(self):
         """Overridden get_queryset method"""
-        return EstablishmentMixinView.get_queryset(self) \
-                                     .has_location() \
-                                     .similar_artisans_producers(slug=self.kwargs.get('slug'))
+        qs = super(ArtisanProducerSimilarListView, self).get_queryset()
+        base_establishment = self.get_base_object()
+
+        if base_establishment:
+            return qs.similar_artisans_producers(base_establishment)
+        else:
+            return qs.none()
 
 
 class EstablishmentTypeListView(generics.ListAPIView):