added .disable_inheritance method for the simple fetch exclusives classes

This commit is contained in:
Wilson Júnior 2013-05-09 13:18:58 -03:00
parent 9dfee83e68
commit 96a964a183
2 changed files with 29 additions and 0 deletions

View File

@ -520,6 +520,15 @@ class QuerySet(object):
queryset._none = True
return queryset
def disable_inheritance(self):
"""
Disable inheritance query, fetch only objects for the query class
"""
if self._document._meta.get('allow_inheritance') is True:
self._initial_query = {"_cls": self._document._class_name}
return self
def clone(self):
"""Creates a copy of the current
:class:`~mongoengine.queryset.QuerySet`

View File

@ -3272,5 +3272,25 @@ class QuerySetTest(unittest.TestCase):
self.assertEqual(outer_count, 7) # outer loop should be executed seven times total
self.assertEqual(inner_total_count, 7 * 7) # inner loop should be executed fourtynine times total
def test_disable_inheritance_queryset(self):
class A(Document):
x = IntField()
y = IntField()
meta = {'allow_inheritance': True}
class B(A):
z = IntField()
A.drop_collection()
A(x=10, y=20).save()
A(x=15, y=30).save()
B(x=20, y=40).save()
B(x=30, y=50).save()
for obj in A.objects.disable_inheritance():
self.assertEqual(obj.__class__, A)
if __name__ == '__main__':
unittest.main()