From 9d12dbad704d8d273684397d642bdaefd2546789 Mon Sep 17 00:00:00 2001 From: Harry Marr Date: Sat, 19 Dec 2009 02:33:01 +0000 Subject: [PATCH] Made _cls etc optional, merged sort to order_by --- mongoengine/base.py | 27 +++++++++++++++++--- mongoengine/queryset.py | 55 ++++++++++++++--------------------------- tests/document.py | 30 ++++++++++++++++++++++ tests/queryset.py | 12 ++++----- 4 files changed, 78 insertions(+), 46 deletions(-) diff --git a/mongoengine/base.py b/mongoengine/base.py index ff6fbd66..60409127 100644 --- a/mongoengine/base.py +++ b/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 - data['_types'] = self._superclasses.keys() + [self._class_name] + # Only add _cls and _types if allow_inheritance is not False + 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'] diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py index 8bb3510c..4f406b44 100644 --- a/mongoengine/queryset.py +++ b/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: - self._ordering.append(sort_rule) - 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 + def order_by(self, *keys): + """Order 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 diff --git a/tests/document.py b/tests/document.py index fdb195de..bec66ee9 100644 --- a/tests/document.py +++ b/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. """ diff --git a/tests/queryset.py b/tests/queryset.py index 39b14fc1..d40e73b0 100644 --- a/tests/queryset.py +++ b/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): - """Ensure that QuerySets may be sorted. + def test_order_by(self): + """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):