Filter an existing django queryset

django_elasticsearch_dsl/search.py

@@ -1,4 +1,5 @@
 from django.db.models import Case, When
+from django.shortcuts import _get_queryset
 
 from elasticsearch_dsl import Search as DSLSearch
 
@@ -13,12 +14,20 @@ def _clone(self):
         s._model = self._model
         return s
 
-    def to_queryset(self, keep_order=True):
+    def filter_queryset(self, klass, keep_search_order=True):
         """
-        This method return a django queryset from the an elasticsearch result.
-        It cost a query to the sql db.
+        Filter an existing django queryset using the elasticsearch result.
+        It costs a query to the sql db.
+        klass may be a Model, Manager, or QuerySet object.
         """
+
+        qs = _get_queryset(klass)
         s = self
+        if s._model is not qs.model:
+            raise TypeError(
+                'Unexpected queryset model '
+                '(should be: %s, got: %s)' % (s._model, qs.model)
+            )
 
         # Do not query again if the es result is already cached
         if not hasattr(self, '_response'):
@@ -27,13 +36,19 @@ def to_queryset(self, keep_order=True):
             s = s.execute()
 
         pks = [result.meta.id for result in s]
+        qs = qs.filter(pk__in=pks)
 
-        qs = self._model.objects.filter(pk__in=pks)
-
-        if keep_order:
+        if keep_search_order:
             preserved_order = Case(
                 *[When(pk=pk, then=pos) for pos, pk in enumerate(pks)]
             )
             qs = qs.order_by(preserved_order)
 
         return qs
+
+    def to_queryset(self, keep_order=True):
+        """
+        Return a django queryset from the elasticsearch result.
+        It costs a query to the sql db.
+        """
+        return self.filter_queryset(self._model, keep_order)

tests/test_integration.py

@@ -329,6 +329,24 @@ def test_rebuild_command(self):
         result = AdDocument().search().execute()
         self.assertEqual(len(result), 3)
 
+    def test_filter_queryset(self):
+        Ad(title="Nothing that match",  car=self.car1).save()
+
+        qs = AdDocument().search().query(
+            'match', title="Ad number 2").filter_queryset(Ad.objects)
+        self.assertEqual(qs.count(), 2)
+        self.assertEqual(list(qs), [self.ad2, self.ad1])
+
+        qs = AdDocument().search().query(
+            'match', title="Ad number 2"
+        ).filter_queryset(Ad.objects.filter(url="www.ad2.com"))
+        self.assertEqual(qs.count(), 1)
+        self.assertEqual(list(qs), [self.ad2])
+
+        with self.assertRaisesMessage(TypeError, 'Unexpected queryset model'):
+            AdDocument().search().query(
+                'match', title="Ad number 2").filter_queryset(Category.objects)
+
     def test_to_queryset(self):
         Ad(title="Nothing that match",  car=self.car1).save()
         qs = AdDocument().search().query(