From ab5dec31f8fc074bf64fcbe297e570123d07fbee Mon Sep 17 00:00:00 2001 From: Thomas Steinacher Date: Thu, 4 Jul 2013 19:15:08 -0700 Subject: [PATCH] QuerySet.only_classes / QuerySet.exclude_classes: Allow to limit/exclude classes for documents that can be inherited. --- mongoengine/queryset/queryset.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/mongoengine/queryset/queryset.py b/mongoengine/queryset/queryset.py index 3b376db..c27adef 100644 --- a/mongoengine/queryset/queryset.py +++ b/mongoengine/queryset/queryset.py @@ -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`