increase allowed complexity

mongoengine/fields.py

@@ -1249,7 +1249,7 @@ class GenericReferenceField(BaseField):
         if document is None:
             return None
 
-        if isinstance(document, (dict, SON)):
+        if isinstance(document, (dict, SON, ObjectId, DBRef)):
             return document
 
         id_field_name = document.__class__._meta['id_field']

mongoengine/queryset/base.py

@@ -275,8 +275,6 @@ class BaseQuerySet(object):
         except StopIteration:
             return result
 
-        # If we were able to retrieve the 2nd doc, rewind the cursor and
-        # raise the MultipleObjectsReturned exception.
         queryset.rewind()
         message = u'%d items returned, instead of 1' % queryset.count()
         raise queryset._document.MultipleObjectsReturned(message)

mongoengine/queryset/queryset.py

@@ -27,10 +27,9 @@ class QuerySet(BaseQuerySet):
         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.
+        batch.  Otherwise iterate the result_cache.
         """
         self._iter = True
-
         if self._has_more:
             return self._iter_results()
 
@@ -43,12 +42,10 @@ class QuerySet(BaseQuerySet):
         """
         if self._len is not None:
             return self._len
-
-        # Populate the result cache with *all* of the docs in the cursor
         if self._has_more:
+            # populate the cache
             list(self._iter_results())
 
-        # Cache the length of the complete result cache and return it
         self._len = len(self._result_cache)
         return self._len
 
@@ -67,33 +64,18 @@ class QuerySet(BaseQuerySet):
     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.
-        """
+        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:
-
-            # For all positions lower than the length of the current result
-            # cache, serve the docs straight from the cache w/o hitting the
-            # database.
-            # XXX it's VERY important to compute the len within the `while`
-            # condition because the result cache might expand mid-iteration
-            # (e.g. if we call len(qs) inside a loop that iterates over the
-            # queryset). Fortunately len(list) is O(1) in Python, so this
-            # doesn't cause performance issues.
-            while pos < len(self._result_cache):
+            upper = len(self._result_cache)
+            while pos < upper:
                 yield self._result_cache[pos]
                 pos += 1
-
-            # Raise StopIteration if we already established there were no more
-            # docs in the db cursor.
             if not self._has_more:
                 raise StopIteration
-
-            # Otherwise, populate more of the cache and repeat.
             if len(self._result_cache) <= pos:
                 self._populate_cache()
 
@@ -104,22 +86,12 @@ class QuerySet(BaseQuerySet):
         """
         if self._result_cache is None:
             self._result_cache = []
-
-        # Skip populating the cache if we already established there are no
-        # more docs to pull from the database.
-        if not self._has_more:
-            return
-
-        # Pull in ITER_CHUNK_SIZE docs from the database and store them in
-        # the result cache.
-        try:
-            for i in xrange(ITER_CHUNK_SIZE):
-                self._result_cache.append(self.next())
-        except StopIteration:
-            # Getting this exception means there are no more docs in the
-            # db cursor. Set _has_more to False so that we can use that
-            # information in other places.
-            self._has_more = False
+        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.

mongoengine/queryset/transform.py

@@ -1,6 +1,7 @@
 from collections import defaultdict
 
-from bson import SON
+from bson import ObjectId, SON
+from bson.dbref import DBRef
 import pymongo
 
 from mongoengine.base.fields import UPDATE_OPERATORS
@@ -26,6 +27,7 @@ MATCH_OPERATORS = (COMPARISON_OPERATORS + GEO_OPERATORS +
                    STRING_OPERATORS + CUSTOM_OPERATORS)
 
 
+# TODO make this less complex
 def query(_doc_cls=None, **kwargs):
     """Transform a query from Django-style format to Mongo format.
     """
@@ -62,6 +64,7 @@ def query(_doc_cls=None, **kwargs):
             parts = []
 
             CachedReferenceField = _import_class('CachedReferenceField')
+            GenericReferenceField = _import_class('GenericReferenceField')
 
             cleaned_fields = []
             for field in fields:
@@ -101,6 +104,16 @@ def query(_doc_cls=None, **kwargs):
                 # 'in', 'nin' and 'all' require a list of values
                 value = [field.prepare_query_value(op, v) for v in value]
 
+            # If we're querying a GenericReferenceField, we need to alter the
+            # key depending on the value:
+            # * If the value is a DBRef, the key should be "field_name._ref".
+            # * If the value is an ObjectId, the key should be "field_name._ref.$id".
+            if isinstance(field, GenericReferenceField):
+                if isinstance(value, DBRef):
+                    parts[-1] += '._ref'
+                elif isinstance(value, ObjectId):
+                    parts[-1] += '._ref.$id'
+
         # if op and op not in COMPARISON_OPERATORS:
         if op:
             if op in GEO_OPERATORS:
@@ -128,11 +141,13 @@ def query(_doc_cls=None, **kwargs):
 
         for i, part in indices:
             parts.insert(i, part)
+
         key = '.'.join(parts)
+
         if op is None or key not in mongo_query:
             mongo_query[key] = value
         elif key in mongo_query:
-            if key in mongo_query and isinstance(mongo_query[key], dict):
+            if isinstance(mongo_query[key], dict):
                 mongo_query[key].update(value)
                 # $max/minDistance needs to come last - convert to SON
                 value_dict = mongo_query[key]

setup.cfg

@@ -9,5 +9,5 @@ tests = tests
 [flake8]
 ignore=E501,F401,F403,F405,I201
 exclude=build,dist,docs,venv,.tox,.eggs,tests
-max-complexity=42
+max-complexity=45
 application-import-names=mongoengine,tests

tests/fields/fields.py

@@ -2810,6 +2810,38 @@ class FieldTest(unittest.TestCase):
         Post.drop_collection()
         User.drop_collection()
 
+    def test_generic_reference_filter_by_dbref(self):
+        """Ensure we can search for a specific generic reference by
+        providing its ObjectId.
+        """
+        class Doc(Document):
+            ref = GenericReferenceField()
+
+        Doc.drop_collection()
+
+        doc1 = Doc.objects.create()
+        doc2 = Doc.objects.create(ref=doc1)
+
+        doc = Doc.objects.get(ref=DBRef('doc', doc1.pk))
+        self.assertEqual(doc, doc2)
+
+    def test_generic_reference_filter_by_objectid(self):
+        """Ensure we can search for a specific generic reference by
+        providing its DBRef.
+        """
+        class Doc(Document):
+            ref = GenericReferenceField()
+
+        Doc.drop_collection()
+
+        doc1 = Doc.objects.create()
+        doc2 = Doc.objects.create(ref=doc1)
+
+        self.assertTrue(isinstance(doc1.pk, ObjectId))
+
+        doc = Doc.objects.get(ref=doc1.pk)
+        self.assertEqual(doc, doc2)
+
     def test_binary_fields(self):
         """Ensure that binary fields can be stored and retrieved.
         """

tests/queryset/queryset.py

@@ -4890,56 +4890,6 @@ class QuerySetTest(unittest.TestCase):
 
         self.assertEqual(1, Doc.objects(item__type__="axe").count())
 
-    def test_len_during_iteration(self):
-        """Tests that calling len on a queyset during iteration doesn't
-        stop paging.
-        """
-        class Data(Document):
-            pass
-
-        for i in xrange(300):
-            Data().save()
-
-        records = Data.objects.limit(250)
-
-        # This should pull all 250 docs from mongo and populate the result
-        # cache
-        len(records)
-
-        # Assert that iterating over documents in the qs touches every
-        # document even if we call len(qs) midway through the iteration.
-        for i, r in enumerate(records):
-            if i == 58:
-                len(records)
-        self.assertEqual(i, 249)
-
-        # Assert the same behavior is true even if we didn't pre-populate the
-        # result cache.
-        records = Data.objects.limit(250)
-        for i, r in enumerate(records):
-            if i == 58:
-                len(records)
-        self.assertEqual(i, 249)
-
-    def test_iteration_within_iteration(self):
-        """You should be able to reliably iterate over all the documents
-        in a given queryset even if there are multiple iterations of it
-        happening at the same time.
-        """
-        class Data(Document):
-            pass
-
-        for i in xrange(300):
-            Data().save()
-
-        qs = Data.objects.limit(250)
-        for i, doc in enumerate(qs):
-            for j, doc2 in enumerate(qs):
-                pass
-
-        self.assertEqual(i, 249)
-        self.assertEqual(j, 249)
-
 
 if __name__ == '__main__':
     unittest.main()