jyapayne/mongoengine
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

QuerySet.only_classes / QuerySet.exclude_classes: Allow to limit/exclude classes for documents that can be inherited.

Browse source
This commit is contained in:
Thomas Steinacher 2013-07-04 19:15:08 -07:00
commit ab5dec31f8
1 changed files with 28 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

28
mongoengine/queryset/queryset.py
View file

@ -589,6 +589,34 @@ class QuerySet(object):
return self
def only_classes(self, *classes):
doc = self._document
if doc._meta.get('allow_inheritance') is True:
queryset = self.clone()
class_names = [cls._class_name for cls in classes]
allowed_class_names = [name for name in self._document._subclasses if name in class_names]
if len(allowed_class_names) == 1:
queryset._initial_query = {"_cls": allowed_class_names[0]}
else:
queryset._initial_query = {"_cls": {"$in": allowed_class_names}}
return queryset
else:
return self
def exclude_classes(self, *classes):
doc = self._document
if doc._meta.get('allow_inheritance') is True:
queryset = self.clone()
class_names = [cls._class_name for cls in classes]
allowed_class_names = [name for name in self._document._subclasses if name in class_names]
if len(allowed_class_names) == 1:
queryset._initial_query = {"_cls": {"$ne": allowed_class_names[0]}}
else:
queryset._initial_query = {"_cls": {"$nin": allowed_class_names}}
return queryset
else:
return self
def clone(self):
"""Creates a copy of the current
:class:`~mongoengine.queryset.QuerySet`