Added support for using custom QuerySet classes

mongoengine/base.py

@@ -255,6 +255,7 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
             'index_background': False,
             'index_drop_dups': False,
             'index_opts': {},
+            'queryset_class': QuerySet,
         }
         meta.update(base_meta)
 

mongoengine/queryset.py

@@ -992,7 +992,8 @@ class QuerySetManager(object):
                 self._collection = db[collection]
 
         # owner is the document that contains the QuerySetManager
-        queryset = QuerySet(owner, self._collection)
+        queryset_class = owner._meta['queryset_class'] or QuerySet
+        queryset = queryset_class(owner, self._collection)
         if self._manager_func:
             if self._manager_func.func_code.co_argcount == 1:
                 queryset = self._manager_func(queryset)

tests/queryset.py

@@ -1378,6 +1378,26 @@ class QTest(unittest.TestCase):
         self.assertEqual(Post.objects.filter(created_user=user).count(), 1)
         self.assertEqual(Post.objects.filter(Q(created_user=user)).count(), 1)
 
+    def test_custom_querysets(self):
+        """Ensure that custom QuerySet classes may be used.
+        """
+        class CustomQuerySet(QuerySet):
+            def not_empty(self):
+                return len(self) > 0
+
+        class Post(Document):
+            meta = {'queryset_class': CustomQuerySet}
+
+        Post.drop_collection()
+
+        self.assertTrue(isinstance(Post.objects, CustomQuerySet))
+        self.assertFalse(Post.objects.not_empty())
+
+        Post().save()
+        self.assertTrue(Post.objects.not_empty())
+
+        Post.drop_collection()
+
 
 if __name__ == '__main__':
     unittest.main()