From 056c604dc30ce616212f22514691d129ba328fca Mon Sep 17 00:00:00 2001
From: Ross Lawley <ross.lawley@gmail.com>
Date: Fri, 22 Jun 2012 16:22:27 +0100
Subject: [PATCH] Fixes __repr__ modifying the cursor

Fixes MongoEngine/mongoengine#30
---
 docs/changelog.rst      |  5 +++++
 mongoengine/queryset.py | 32 +++++++++++++++++++-------------
 tests/test_queryset.py  | 39 ++++++++++++++++++++++++++++++---------
 3 files changed, 54 insertions(+), 22 deletions(-)

diff --git a/docs/changelog.rst b/docs/changelog.rst
index fdf7c7bc..c5489334 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -2,6 +2,11 @@
 Changelog
 =========
 
+Changes in 0.6.X
+================
+
+- Fixes __repr__ modifying the cursor
+
 Changes in 0.6.12
 =================
 - Fixes scalar lookups for primary_key
diff --git a/mongoengine/queryset.py b/mongoengine/queryset.py
index 3ebbb87e..21e3f9e5 100644
--- a/mongoengine/queryset.py
+++ b/mongoengine/queryset.py
@@ -341,6 +341,7 @@ class QuerySet(object):
         self._timeout = True
         self._class_check = True
         self._slave_okay = False
+        self._iter = False
         self._scalar = []
 
         # If inheritance is allowed, only return instances and instances of
@@ -953,6 +954,7 @@ class QuerySet(object):
     def next(self):
         """Wrap the result in a :class:`~mongoengine.Document` object.
         """
+        self._iter = True
         try:
             if self._limit == 0:
                 raise StopIteration
@@ -969,6 +971,7 @@ class QuerySet(object):
 
         .. versionadded:: 0.3
         """
+        self._iter = False
         self._cursor.rewind()
 
     def count(self):
@@ -1808,21 +1811,24 @@ class QuerySet(object):
         return data
 
     def __repr__(self):
-        limit = REPR_OUTPUT_SIZE + 1
-        start = (0 if self._skip is None else self._skip)
-        if self._limit is None:
-            stop = start + limit
-        if self._limit is not None:
-            if self._limit - start > limit:
-                stop = start + limit
-            else:
-                stop = self._limit
-        try:
-            data = list(self[start:stop])
-        except pymongo.errors.InvalidOperation:
-            return ".. queryset mid-iteration .."
+        """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 select_related(self, max_depth=1):
diff --git a/tests/test_queryset.py b/tests/test_queryset.py
index aa6eabb5..e7ee6780 100644
--- a/tests/test_queryset.py
+++ b/tests/test_queryset.py
@@ -636,17 +636,38 @@ class QuerySetTest(unittest.TestCase):
         self.assertEqual(people1, people2)
         self.assertEqual(people1, people3)
 
-    def test_repr_iteration(self):
-        """Ensure that QuerySet __repr__ can handle loops
-        """
-        self.Person(name='Person 1').save()
-        self.Person(name='Person 2').save()
+    def test_repr(self):
+        """Test repr behavior isnt destructive"""
 
-        queryset = self.Person.objects
-        self.assertEquals('[<Person: Person object>, <Person: Person object>]', repr(queryset))
-        for person in queryset:
-            self.assertEquals('.. queryset mid-iteration ..', repr(queryset))
+        class Doc(Document):
+            number = IntField()
 
+            def __repr__(self):
+               return "<Doc: %s>" % self.number
+
+        Doc.drop_collection()
+
+        for i in xrange(1000):
+            Doc(number=i).save()
+
+        docs = Doc.objects.order_by('number')
+
+        self.assertEquals(docs.count(), 1000)
+        self.assertEquals(len(docs), 1000)
+
+        docs_string = "%s" % docs
+        self.assertTrue("Doc: 0" in docs_string)
+
+        self.assertEquals(docs.count(), 1000)
+        self.assertEquals(len(docs), 1000)
+
+        # Limit and skip
+        self.assertEquals('[<Doc: 1>, <Doc: 2>, <Doc: 3>]', "%s" % docs[1:4])
+
+        self.assertEquals(docs.count(), 3)
+        self.assertEquals(len(docs), 3)
+        for doc in docs:
+            self.assertEqual('.. queryset mid-iteration ..', repr(docs))
 
     def test_regex_query_shortcuts(self):
         """Ensure that contains, startswith, endswith, etc work.