fix slicing QuerySet.count()

Author: timgraham

django_mongodb/compiler.py

@@ -126,7 +126,17 @@ def get_count(self, check_exists=False):
 
         If `check_exists` is True, only check if any object matches.
         """
-        kwargs = {"limit": 1} if check_exists else {}
+        kwargs = {}
+        # If this query is sliced, the limits will be set on the subquery.
+        inner_query = getattr(self.query, "inner_query", None)
+        low_mark = inner_query.low_mark if inner_query else 0
+        high_mark = inner_query.high_mark if inner_query else None
+        if low_mark > 0:
+            kwargs["skip"] = low_mark
+        if check_exists:
+            kwargs["limit"] = 1
+        elif high_mark is not None:
+            kwargs["limit"] = high_mark - low_mark
         try:
             return self.build_query().count(**kwargs)
         except EmptyResultSet:

django_mongodb/features.py

@@ -26,8 +26,6 @@ class DatabaseFeatures(BaseDatabaseFeatures):
         "lookup.tests.LookupTests.test_exact_none_transform",
         # "Save with update_fields did not affect any rows."
         "basic.tests.SelectOnSaveTests.test_select_on_save_lying_update",
-        # Slicing with QuerySet.count() doesn't work.
-        "lookup.tests.LookupTests.test_count",
         # Lookup in order_by() not supported:
         # argument of type '<database function>' is not iterable
         "db_functions.comparison.test_coalesce.CoalesceTests.test_ordering",

django_mongodb/query.py

@@ -51,12 +51,16 @@ def fetch(self):
         yield from self.get_cursor()
 
     @wrap_database_errors
-    def count(self, limit=None):
+    def count(self, limit=None, skip=None):
         """
         Return the number of objects that would be returned, if this query was
-        executed, up to `limit`.
+        executed, up to `limit`, skipping `skip`.
         """
-        kwargs = {"limit": limit} if limit is not None else {}
+        kwargs = {}
+        if limit is not None:
+            kwargs["limit"] = limit
+        if skip is not None:
+            kwargs["skip"] = skip
         return self.collection.count_documents(self.mongo_query, **kwargs)
 
     def order_by(self, ordering):