Made _cls etc optional, merged sort to order_by

mongoengine/base.py

@@ -136,17 +136,32 @@ class TopLevelDocumentMetaclass(DocumentMetaclass):
         if attrs.get('__metaclass__') == TopLevelDocumentMetaclass:
             return super_new(cls, name, bases, attrs)
 
-        collection = attrs.get('__collection__', name.lower())
+        collection = name.lower()
         
+        simple_class = True
         # Subclassed documents inherit collection from superclass
         for base in bases:
             if hasattr(base, '_meta') and 'collection' in base._meta:
+                # Ensure that the Document class may be subclassed - 
+                # inheritance may be disabled to remove dependency on 
+                # additional fields _cls and _types
+                if base._meta.get('allow_inheritance', True) == False:
+                    raise ValueError('Document %s may not be subclassed' %
+                                     base.__name__)
+                else:
+                    simple_class = False
                 collection = base._meta['collection']
 
         meta = {
             'collection': collection,
+            'allow_inheritance': True,
         }
         meta.update(attrs.get('meta', {}))
+        # Only simple classes - direct subclasses of Document - may set
+        # allow_inheritance to False
+        if not simple_class and not meta['allow_inheritance']:
+            raise ValueError('Only direct subclasses of Document may set '
+                             '"allow_inheritance" to False')
         attrs['_meta'] = meta
 
         attrs['id'] = ObjectIdField(name='_id')
@@ -228,8 +243,11 @@ class BaseDocument(object):
             value = getattr(self, field_name, None)
             if value is not None:
                 data[field.name] = field.to_mongo(value)
-        data['_cls'] = self._class_name
+        # Only add _cls and _types if allow_inheritance is not False
-        data['_types'] = self._superclasses.keys() + [self._class_name]
+        if not (hasattr(self, '_meta') and
+                self._meta.get('allow_inheritance', True) == False):
+            data['_cls'] = self._class_name
+            data['_types'] = self._superclasses.keys() + [self._class_name]
         return data
     
     @classmethod
@@ -242,6 +260,9 @@ class BaseDocument(object):
 
         data = dict((str(key), value) for key, value in son.items())
 
+        if '_types' in data:
+            del data['_types']
+
         if '_cls' in data:
             del data['_cls']
 

mongoengine/queryset.py

@@ -12,8 +12,11 @@ class QuerySet(object):
         self._document = document
         self._collection = collection
         self._query = {}
+        # If inheritance is allowed, only return instances and instances of
+        # subclasses of the class being used
+        if document._meta.get('allow_inheritance'):
+            self._query = {'_types': self._document._class_name}
         self._cursor_obj = None
-        self._ordering = []
         
     def ensure_index(self, key_or_list, direction=None):
         """Ensure that the given indexes are in place.
@@ -105,43 +108,9 @@ class QuerySet(object):
         """
         self._cursor.skip(n)
         return self
-        
-    def order_by(self, *params):
-        """Apply ordering conditions, Django-style.
-        
-        e.g., ``Model.objects.().order_by("-published_date", "ordering")``
-        will order first by ``published_date DESC``, and then ``ordering ASC``.
-        
-        """
-        if not params:
-            self._ordering = []
-        for param in params:
-            if param.startswith("-"):
-                param = param[1:]
-                sort_dir = pymongo.DESCENDING
-            else:
-                sort_dir = pymongo.ASCENDING
-            sort_rule = (param, sort_dir)
 
-            if not sort_rule in self._ordering:
+    def order_by(self, *keys):
-                self._ordering.append(sort_rule)
+        """Order the QuerySet by the keys. The order may be specified by 
-        self._cursor.sort(self._ordering)
-        return self
-        
-    def explain(self, format=False):
-        plan = self._cursor.explain()
-        if format:
-            import pprint
-            plan = pprint.pformat(plan)
-        return plan
-        
-    def delete(self):
-        """Delete the documents matched by the query.
-        """
-        self._collection.remove(self._query)
-
-    def sort(self, *keys):
-        """Sort the QuerySet by the keys. The order may be specified by 
         prepending each of the keys by a + or a -. Ascending order is assumed.
         """
         key_list = []
@@ -155,6 +124,18 @@ class QuerySet(object):
 
         self._cursor.sort(key_list)
         return self
+        
+    def explain(self, format=False):
+        plan = self._cursor.explain()
+        if format:
+            import pprint
+            plan = pprint.pformat(plan)
+        return plan
+        
+    def delete(self):
+        """Delete the documents matched by the query.
+        """
+        self._collection.remove(self._query)
 
     def __iter__(self):
         return self

tests/document.py

@@ -126,6 +126,36 @@ class DocumentTest(unittest.TestCase):
         self.assertEqual(Employee._meta['collection'], 
                          self.Person._meta['collection'])
 
+    def test_allow_inheritance(self):
+        """Ensure that inheritance may be disabled on simple classes and that
+        _cls and _types will not be used.
+        """
+        class Animal(Document):
+            meta = {'allow_inheritance': False}
+            name = StringField()
+
+        Animal.drop_collection()
+
+        def create_dog_class():
+            class Dog(Animal):
+                pass
+        self.assertRaises(ValueError, create_dog_class)
+        
+        # Check that _cls etc aren't present on simple documents
+        dog = Animal(name='dog')
+        dog.save()
+        collection = self.db[Animal._meta['collection']]
+        obj = collection.find_one()
+        self.assertFalse('_cls' in obj)
+        self.assertFalse('_types' in obj)
+
+        Animal.drop_collection()
+
+        def create_employee_class():
+            class Employee(self.Person):
+                meta = {'allow_inheritance': False}
+        self.assertRaises(ValueError, create_employee_class)
+
     def test_creation(self):
         """Ensure that document may be created using keyword arguments.
         """

tests/queryset.py

@@ -137,23 +137,23 @@ class QuerySetTest(unittest.TestCase):
         self.Person.objects.delete()
         self.assertEqual(self.Person.objects.count(), 0)
 
-    def test_sort(self):
+    def test_order_by(self):
-        """Ensure that QuerySets may be sorted.
+        """Ensure that QuerySets may be ordered.
         """
         self.Person(name="User A", age=20).save()
         self.Person(name="User B", age=40).save()
         self.Person(name="User C", age=30).save()
 
-        names = [p.name for p in self.Person.objects.sort('-age')]
+        names = [p.name for p in self.Person.objects.order_by('-age')]
         self.assertEqual(names, ['User B', 'User C', 'User A'])
 
-        names = [p.name for p in self.Person.objects.sort('+age')]
+        names = [p.name for p in self.Person.objects.order_by('+age')]
         self.assertEqual(names, ['User A', 'User C', 'User B'])
 
-        names = [p.name for p in self.Person.objects.sort('age')]
+        names = [p.name for p in self.Person.objects.order_by('age')]
         self.assertEqual(names, ['User A', 'User C', 'User B'])
         
-        ages = [p.age for p in self.Person.objects.sort('-name')]
+        ages = [p.age for p in self.Person.objects.order_by('-name')]
         self.assertEqual(ages, [30, 40, 20])
 
     def tearDown(self):