Added no_sub_classes context manager and queryset helper (#312)

mongoengine/context_managers.py

@@ -1,8 +1,10 @@
 from mongoengine.common import _import_class
 from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
-from mongoengine.queryset import OperationError, QuerySet
+from mongoengine.queryset import QuerySet
 
-__all__ = ("switch_db", "switch_collection", "no_dereference", "query_counter")
+
+__all__ = ("switch_db", "switch_collection", "no_dereference",
+           "no_sub_classes", "query_counter")
 
 
 class switch_db(object):
@@ -130,6 +132,36 @@ class no_dereference(object):
         return self.cls
 
 
+class no_sub_classes(object):
+    """ no_sub_classes context manager.
+
+    Only returns instances of this class and no sub (inherited) classes::
+
+        with no_sub_classes(Group) as Group:
+            Group.objects.find()
+
+    """
+
+    def __init__(self, cls):
+        """ Construct the no_sub_classes context manager.
+
+        :param cls: the class to turn querying sub classes on
+        """
+        self.cls = cls
+
+    def __enter__(self):
+        """ change the objects default and _auto_dereference values"""
+        self.cls._all_subclasses = self.cls._subclasses
+        self.cls._subclasses = (self.cls,)
+        return self.cls
+
+    def __exit__(self, t, value, traceback):
+        """ Reset the default and _auto_dereference values"""
+        self.cls._subclasses = self.cls._all_subclasses
+        delattr(self.cls, '_all_subclasses')
+        return self.cls
+
+
 class QuerySetNoDeRef(QuerySet):
     """Special no_dereference QuerySet"""
     def __dereference(items, max_depth=1, instance=None, name=None):

mongoengine/queryset/queryset.py

@@ -569,9 +569,9 @@ class QuerySet(object):
         queryset._none = True
         return queryset
 
-    def disable_inheritance(self):
+    def no_sub_classes(self):
         """
-        Disable inheritance query, fetch only objects for the query class
+        Only return instances of this document and not any inherited documents
         """
         if self._document._meta.get('allow_inheritance') is True:
             self._initial_query = {"_cls": self._document._class_name}