diff --git a/docs/changelog.rst b/docs/changelog.rst index 8fc279e7..9e1cec80 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -21,6 +21,7 @@ Changes in 0.8 - Remove _types and just use _cls for inheritance (#148) - Only allow QNode instances to be passed as query objects (#199) - Dynamic fields are now validated on save (#153) (#154) +- Added support for multiple slices and made slicing chainable. (#170) (#190) (#191) Changes in 0.7.9 ================ diff --git a/mongoengine/queryset/field_list.py b/mongoengine/queryset/field_list.py index 1c825fa9..7b2b0cb5 100644 --- a/mongoengine/queryset/field_list.py +++ b/mongoengine/queryset/field_list.py @@ -12,20 +12,31 @@ class QueryFieldList(object): self.fields = set(fields) self.always_include = set(always_include) self._id = None + self.slice = {} def __add__(self, f): - if not self.fields: + if isinstance(f.value, dict): + for field in f.fields: + self.slice[field] = f.value + if not self.fields: + self.fields = f.fields + elif not self.fields: self.fields = f.fields self.value = f.value + self.slice = {} elif self.value is self.ONLY and f.value is self.ONLY: + self._clean_slice() self.fields = self.fields.intersection(f.fields) elif self.value is self.EXCLUDE and f.value is self.EXCLUDE: self.fields = self.fields.union(f.fields) + self._clean_slice() elif self.value is self.ONLY and f.value is self.EXCLUDE: self.fields -= f.fields + self._clean_slice() elif self.value is self.EXCLUDE and f.value is self.ONLY: self.value = self.ONLY self.fields = f.fields - self.fields + self._clean_slice() if '_id' in f.fields: self._id = f.value @@ -42,10 +53,18 @@ class QueryFieldList(object): def as_dict(self): field_list = dict((field, self.value) for field in self.fields) + if self.slice: + field_list.update(self.slice) if self._id is not None: field_list['_id'] = self._id return field_list def reset(self): self.fields = set([]) + self.slice = {} self.value = self.ONLY + + def _clean_slice(self): + if self.slice: + for field in set(self.slice.keys()) - self.fields: + del self.slice[field] diff --git a/tests/queryset/field_list.py b/tests/queryset/field_list.py index 9e711336..4a8a72b3 100644 --- a/tests/queryset/field_list.py +++ b/tests/queryset/field_list.py @@ -115,6 +115,35 @@ class OnlyExcludeAllTest(unittest.TestCase): qs = qs.only(*only) self.assertEqual(qs._loaded_fields.as_dict(), {'b': 1, 'c': 1}) + def test_slicing(self): + + class MyDoc(Document): + a = ListField() + b = ListField() + c = ListField() + d = ListField() + e = ListField() + f = ListField() + + include = ['a', 'b', 'c', 'd', 'e'] + exclude = ['d', 'e'] + only = ['b', 'c'] + + qs = MyDoc.objects.fields(**dict(((i, 1) for i in include))) + qs = qs.exclude(*exclude) + qs = qs.only(*only) + qs = qs.fields(slice__b=5) + self.assertEqual(qs._loaded_fields.as_dict(), + {'b': {'$slice': 5}, 'c': 1}) + + qs = qs.fields(slice__c=[5, 1]) + self.assertEqual(qs._loaded_fields.as_dict(), + {'b': {'$slice': 5}, 'c': {'$slice': [5, 1]}}) + + qs = qs.exclude('c') + self.assertEqual(qs._loaded_fields.as_dict(), + {'b': {'$slice': 5}}) + def test_only(self): """Ensure that QuerySet.only only returns the requested fields. """