From b887ea96236c62a04a6962dabc747dda14eb1057 Mon Sep 17 00:00:00 2001 From: otrofimov Date: Thu, 8 Aug 2019 11:55:45 +0300 Subject: [PATCH 1/3] Implement collation for queryset --- mongoengine/queryset/base.py | 31 +++++++++++++++++++++++++++++++ tests/document/indexes.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 60 insertions(+) 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. """ From fbb3bf869c9cdea0b6c5060e2f57fabf5b8c5e5d Mon Sep 17 00:00:00 2001 From: otrofimov Date: Thu, 8 Aug 2019 15:56:20 +0300 Subject: [PATCH 2/3] compatibility with black --- tests/document/indexes.py | 27 ++++++++++++++------------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/tests/document/indexes.py b/tests/document/indexes.py index 0bc23d1c..fa3d1706 100644 --- a/tests/document/indexes.py +++ b/tests/document/indexes.py @@ -540,33 +540,34 @@ class IndexesTest(unittest.TestCase): BlogPost.objects.hint(("tags", 1)).count() def test_collation(self): - base = {'locale': "en", 'strength': 2} + base = {"locale": "en", "strength": 2} class BlogPost(Document): name = StringField() - meta = {"indexes": [ - {"fields": ["name"], "name": 'name_index', - 'collation': base} - ]} + 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)) + 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())) + 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'} + 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)) + 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. From 71a6f3d1a46702bf5d66e704519671e828626cd2 Mon Sep 17 00:00:00 2001 From: otrofimov Date: Wed, 21 Aug 2019 18:26:10 +0300 Subject: [PATCH 3/3] test_collation: Added test with `pymongo.collation.Collation` object Readable list of BlogPost names for test --- tests/document/indexes.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/tests/document/indexes.py b/tests/document/indexes.py index fa3d1706..dcd3fc6a 100644 --- a/tests/document/indexes.py +++ b/tests/document/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 import pymongo from six import iteritems @@ -552,7 +553,7 @@ class IndexesTest(unittest.TestCase): BlogPost.drop_collection() - names = tuple("%sag %i" % ("t" if n % 2 == 0 else "T", n) for n in range(10)) + names = ["tag1", "Tag2", "tag3", "Tag4", "tag5"] for name in names: BlogPost(name=name).save() @@ -560,7 +561,13 @@ class IndexesTest(unittest.TestCase): self.assertEqual( [x.name for x in query_result], sorted(names, key=lambda x: x.lower()) ) - self.assertEqual(10, query_result.count()) + 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):