164 lines
5.0 KiB
Python
164 lines
5.0 KiB
Python
from mongoengine.errors import OperationError
|
|
from mongoengine.queryset.base import (BaseQuerySet, DO_NOTHING, NULLIFY,
|
|
CASCADE, DENY, PULL)
|
|
|
|
__all__ = ('QuerySet', 'QuerySetNoCache', 'DO_NOTHING', 'NULLIFY', 'CASCADE',
|
|
'DENY', 'PULL')
|
|
|
|
# The maximum number of items to display in a QuerySet.__repr__
|
|
REPR_OUTPUT_SIZE = 20
|
|
ITER_CHUNK_SIZE = 100
|
|
|
|
|
|
class QuerySet(BaseQuerySet):
|
|
"""The default queryset, that builds queries and handles a set of results
|
|
returned from a query.
|
|
|
|
Wraps a MongoDB cursor, providing :class:`~mongoengine.Document` objects as
|
|
the results.
|
|
"""
|
|
|
|
_has_more = True
|
|
_len = None
|
|
_result_cache = None
|
|
|
|
def __iter__(self):
|
|
"""Iteration utilises a results cache which iterates the cursor
|
|
in batches of ``ITER_CHUNK_SIZE``.
|
|
|
|
If ``self._has_more`` the cursor hasn't been exhausted so cache then
|
|
batch. Otherwise iterate the result_cache.
|
|
"""
|
|
self._iter = True
|
|
if self._has_more:
|
|
return self._iter_results()
|
|
|
|
# iterating over the cache.
|
|
return iter(self._result_cache)
|
|
|
|
def __len__(self):
|
|
"""Since __len__ is called quite frequently (for example, as part of
|
|
list(qs)), we populate the result cache and cache the length.
|
|
"""
|
|
if self._len is not None:
|
|
return self._len
|
|
if self._has_more:
|
|
# populate the cache
|
|
list(self._iter_results())
|
|
|
|
self._len = len(self._result_cache)
|
|
return self._len
|
|
|
|
def __repr__(self):
|
|
"""Provides the string representation of the QuerySet
|
|
"""
|
|
if self._iter:
|
|
return '.. queryset mid-iteration ..'
|
|
|
|
self._populate_cache()
|
|
data = self._result_cache[:REPR_OUTPUT_SIZE + 1]
|
|
if len(data) > REPR_OUTPUT_SIZE:
|
|
data[-1] = "...(remaining elements truncated)..."
|
|
return repr(data)
|
|
|
|
def _iter_results(self):
|
|
"""A generator for iterating over the result cache.
|
|
|
|
Also populates the cache if there are more possible results to yield.
|
|
Raises StopIteration when there are no more results"""
|
|
if self._result_cache is None:
|
|
self._result_cache = []
|
|
pos = 0
|
|
while True:
|
|
upper = len(self._result_cache)
|
|
while pos < upper:
|
|
yield self._result_cache[pos]
|
|
pos += 1
|
|
if not self._has_more:
|
|
raise StopIteration
|
|
if len(self._result_cache) <= pos:
|
|
self._populate_cache()
|
|
|
|
def _populate_cache(self):
|
|
"""
|
|
Populates the result cache with ``ITER_CHUNK_SIZE`` more entries
|
|
(until the cursor is exhausted).
|
|
"""
|
|
if self._result_cache is None:
|
|
self._result_cache = []
|
|
if self._has_more:
|
|
try:
|
|
for i in xrange(ITER_CHUNK_SIZE):
|
|
self._result_cache.append(self.next())
|
|
except StopIteration:
|
|
self._has_more = False
|
|
|
|
def count(self, with_limit_and_skip=False):
|
|
"""Count the selected elements in the query.
|
|
|
|
:param with_limit_and_skip (optional): take any :meth:`limit` or
|
|
:meth:`skip` that has been applied to this cursor into account when
|
|
getting the count
|
|
"""
|
|
if with_limit_and_skip is False:
|
|
return super(QuerySet, self).count(with_limit_and_skip)
|
|
|
|
if self._len is None:
|
|
self._len = super(QuerySet, self).count(with_limit_and_skip)
|
|
|
|
return self._len
|
|
|
|
def no_cache(self):
|
|
"""Convert to a non_caching queryset
|
|
|
|
.. versionadded:: 0.8.3 Convert to non caching queryset
|
|
"""
|
|
if self._result_cache is not None:
|
|
raise OperationError("QuerySet already cached")
|
|
return self.clone_into(QuerySetNoCache(self._document, self._collection))
|
|
|
|
|
|
class QuerySetNoCache(BaseQuerySet):
|
|
"""A non caching QuerySet"""
|
|
|
|
def cache(self):
|
|
"""Convert to a caching queryset
|
|
|
|
.. versionadded:: 0.8.3 Convert to caching queryset
|
|
"""
|
|
return self.clone_into(QuerySet(self._document, self._collection))
|
|
|
|
def __repr__(self):
|
|
"""Provides the string representation of the QuerySet
|
|
|
|
.. versionchanged:: 0.6.13 Now doesnt modify the cursor
|
|
"""
|
|
if self._iter:
|
|
return '.. queryset mid-iteration ..'
|
|
|
|
data = []
|
|
for i in xrange(REPR_OUTPUT_SIZE + 1):
|
|
try:
|
|
data.append(self.next())
|
|
except StopIteration:
|
|
break
|
|
if len(data) > REPR_OUTPUT_SIZE:
|
|
data[-1] = "...(remaining elements truncated)..."
|
|
|
|
self.rewind()
|
|
return repr(data)
|
|
|
|
def __iter__(self):
|
|
queryset = self
|
|
if queryset._iter:
|
|
queryset = self.clone()
|
|
queryset.rewind()
|
|
return queryset
|
|
|
|
|
|
class QuerySetNoDeRef(QuerySet):
|
|
"""Special no_dereference QuerySet"""
|
|
|
|
def __dereference(items, max_depth=1, instance=None, name=None):
|
|
return items
|