Merge pull request #1969 from buggyspace/master

Fix .only() working improperly after using .count() of the same instance of QuerySet

AUTHORS

@@ -248,3 +248,4 @@ that much better:
  * Andy Yankovsky (https://github.com/werat)
  * Bastien Gérard (https://github.com/bagerard)
  * Trevor Hall (https://github.com/tjhall13)
+ * Gleb Voropaev (https://github.com/buggyspace)

docs/changelog.rst

@@ -5,6 +5,7 @@ Changelog
 Development
 ===========
 - (Fill this out as you fix issues and develop your features).
+- Fix .only() working improperly after using .count() of the same instance of QuerySet
 
 =================
 Changes in 0.16.3

mongoengine/queryset/base.py

@@ -396,7 +396,9 @@ class BaseQuerySet(object):
         """
         if self._limit == 0 and with_limit_and_skip is False or self._none:
             return 0
-        return self._cursor.count(with_limit_and_skip=with_limit_and_skip)
+        count = self._cursor.count(with_limit_and_skip=with_limit_and_skip)
+        self._cursor_obj = None
+        return count
 
     def delete(self, write_concern=None, _from_doc_delete=False,
                cascade_refs=None):

tests/queryset/queryset.py

@@ -4720,6 +4720,28 @@ class QuerySetTest(unittest.TestCase):
             'password_salt').only('email').to_json()
         self.assertEqual('[{"email": "ross@example.com"}]', serialized_user)
 
+    def test_only_after_count(self):
+        """Test that only() works after count()"""
+
+        class User(Document):
+            name = StringField()
+            age = IntField()
+            address = StringField()
+        User.drop_collection()
+        User(name="User", age=50,
+             address="Moscow, Russia").save()
+
+        user_queryset = User.objects(age=50)
+
+        result = user_queryset.only("name", "age").as_pymongo().first()
+        self.assertEqual(result, {"name": "User", "age": 50})
+
+        result = user_queryset.count()
+        self.assertEqual(result, 1)
+
+        result = user_queryset.only("name", "age").as_pymongo().first()
+        self.assertEqual(result, {"name": "User", "age": 50})
+
     def test_no_dereference(self):
 
         class Organization(Document):