diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index 9dfb0f19..cde06d54 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -81,6 +81,7 @@ class BaseQuerySet(object): self._limit = None self._skip = None self._hint = -1 # Using -1 as None is a valid value for hint + self._collation = None self._batch_size = None self.only_fields = [] self._max_time_ms = None @@ -782,6 +783,7 @@ class BaseQuerySet(object): "_limit", "_skip", "_hint", + "_collation", "_auto_dereference", "_search_text", "only_fields", @@ -864,6 +866,32 @@ class BaseQuerySet(object): return queryset + def collation(self, collation=None): + """ + Collation allows users to specify language-specific rules for string + comparison, such as rules for lettercase and accent marks. + :param collation: `~pymongo.collation.Collation` or dict with + following fields: + { + locale: str, + caseLevel: bool, + caseFirst: str, + strength: int, + numericOrdering: bool, + alternate: str, + maxVariable: str, + backwards: str + } + Collation should be added to indexes like in test example + """ + queryset = self.clone() + queryset._collation = collation + + if queryset._cursor_obj: + queryset._cursor_obj.collation(collation) + + return queryset + def batch_size(self, size): """Limit the number of documents returned in a single batch (each batch requires a round trip to the server). @@ -1640,6 +1668,9 @@ class BaseQuerySet(object): if self._hint != -1: self._cursor_obj.hint(self._hint) + if self._collation is not None: + self._cursor_obj.collation(self._collation) + if self._batch_size is not None: self._cursor_obj.batch_size(self._batch_size) diff --git a/tests/document/test_indexes.py b/tests/document/test_indexes.py index f94eb359..1b0304c4 100644 --- a/tests/document/test_indexes.py +++ b/tests/document/test_indexes.py @@ -3,6 +3,7 @@ import unittest from datetime import datetime from nose.plugins.skip import SkipTest +from pymongo.collation import Collation from pymongo.errors import OperationFailure from six import iteritems @@ -536,6 +537,42 @@ class TestIndexes(unittest.TestCase): with self.assertRaises(ValueError): BlogPost.objects.hint(("tags", 1)).count() + def test_collation(self): + base = {"locale": "en", "strength": 2} + + class BlogPost(Document): + name = StringField() + meta = { + "indexes": [ + {"fields": ["name"], "name": "name_index", "collation": base} + ] + } + + BlogPost.drop_collection() + + names = ["tag1", "Tag2", "tag3", "Tag4", "tag5"] + for name in names: + BlogPost(name=name).save() + + query_result = BlogPost.objects.collation(base).order_by("name") + self.assertEqual( + [x.name for x in query_result], sorted(names, key=lambda x: x.lower()) + ) + self.assertEqual(5, query_result.count()) + + query_result = BlogPost.objects.collation(Collation(**base)).order_by("name") + self.assertEqual( + [x.name for x in query_result], sorted(names, key=lambda x: x.lower()) + ) + self.assertEqual(5, query_result.count()) + + incorrect_collation = {"arndom": "wrdo"} + with self.assertRaises(OperationFailure): + BlogPost.objects.collation(incorrect_collation).count() + + query_result = BlogPost.objects.collation({}).order_by("name") + self.assertEqual([x.name for x in query_result], sorted(names)) + def test_unique(self): """Ensure that uniqueness constraints are applied to fields. """