Avoid to open all documents from cursors in an if stmt

Using a cursos in an if statement:

cursor = Collection.objects

	if cursor:
		(...)

Will open all documents, because there are not an __nonzero__ method.
This change check only one document (if present) and returns True or False.
This commit is contained in:
Jonathan Prates 2014-05-27 16:33:38 -03:00
parent f099dc6a37
commit 3faf3c84be
2 changed files with 32 additions and 0 deletions

View File

@ -154,6 +154,15 @@ class BaseQuerySet(object):
def __iter__(self):
raise NotImplementedError
def __nonzero__(self):
""" Avoid to open all records in an if stmt """
for value in self:
self.rewind()
return True
return False
# Core functions
def all(self):

View File

@ -3814,6 +3814,29 @@ class QuerySetTest(unittest.TestCase):
self.assertEqual(Example.objects(size=instance_size).count(), 1)
self.assertEqual(Example.objects(size__in=[instance_size]).count(), 1)
def test_cursor_in_an_if_stmt(self):
class Test(Document):
test_field = StringField()
Test.drop_collection()
queryset = Test.objects
if queryset:
raise AssertionError('Empty cursor returns True')
test = Test()
test.test_field = 'test'
test.save()
queryset = Test.objects
if not test:
raise AssertionError('There is data, but cursor returned False')
queryset.next()
if queryset:
raise AssertionError('There is no data left in cursor')
if __name__ == '__main__':
unittest.main()