diff --git a/mongoengine/queryset/base.py b/mongoengine/queryset/base.py index ba3ac95a..b0e1bff2 100644 --- a/mongoengine/queryset/base.py +++ b/mongoengine/queryset/base.py @@ -80,6 +80,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 @@ -781,6 +782,7 @@ class BaseQuerySet(object): "_limit", "_skip", "_hint", + "_collation", "_auto_dereference", "_search_text", "only_fields", @@ -863,6 +865,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). @@ -1636,6 +1664,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/indexes.py b/tests/document/indexes.py index 570e619e..0bc23d1c 100644 --- a/tests/document/indexes.py +++ b/tests/document/indexes.py @@ -539,6 +539,35 @@ class IndexesTest(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 = tuple("%sag %i" % ('t' if n % 2 == 0 else 'T', n) for n in range(10)) + 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(10, 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. """