PYTHON-4549 - Optimize Cursor.to_list (mongodb#1749)

Author: NoahStapp

gridfs/asynchronous/grid_file.py

@@ -1892,6 +1892,9 @@ async def next(self) -> AsyncGridOut:
         next_file = await super().next()
         return AsyncGridOut(self._root_collection, file_document=next_file, session=self.session)
 
+    async def to_list(self) -> list[AsyncGridOut]:
+        return [x async for x in self]  # noqa: C416,RUF100
+
     __anext__ = next
 
     def add_option(self, *args: Any, **kwargs: Any) -> NoReturn:

gridfs/synchronous/grid_file.py

@@ -1878,6 +1878,9 @@ def next(self) -> GridOut:
         next_file = super().next()
         return GridOut(self._root_collection, file_document=next_file, session=self.session)
 
+    def to_list(self) -> list[GridOut]:
+        return [x for x in self]  # noqa: C416,RUF100
+
     __next__ = next
 
     def add_option(self, *args: Any, **kwargs: Any) -> NoReturn:

pymongo/asynchronous/command_cursor.py

@@ -346,6 +346,17 @@ async def _try_next(self, get_more_allowed: bool) -> Optional[_DocumentType]:
         else:
             return None
 
+    async def _next_batch(self, result: list) -> bool:
+        """Get all available documents from the cursor."""
+        if not len(self._data) and not self._killed:
+            await self._refresh()
+        if len(self._data):
+            result.extend(self._data)
+            self._data.clear()
+            return True
+        else:
+            return False
+
     async def try_next(self) -> Optional[_DocumentType]:
         """Advance the cursor without blocking indefinitely.
 
@@ -371,7 +382,11 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         await self.close()
 
     async def to_list(self) -> list[_DocumentType]:
-        return [x async for x in self]  # noqa: C416,RUF100
+        res: list[_DocumentType] = []
+        while self.alive:
+            if not await self._next_batch(res):
+                break
+        return res
 
 
 class AsyncRawBatchCommandCursor(AsyncCommandCursor[_DocumentType]):

pymongo/asynchronous/cursor.py

@@ -1260,6 +1260,20 @@ async def next(self) -> _DocumentType:
         else:
             raise StopAsyncIteration
 
+    async def _next_batch(self, result: list) -> bool:
+        """Get all available documents from the cursor."""
+        if not self._exhaust_checked:
+            self._exhaust_checked = True
+            await self._supports_exhaust()
+        if self._empty:
+            return False
+        if len(self._data) or await self._refresh():
+            result.extend(self._data)
+            self._data.clear()
+            return True
+        else:
+            return False
+
     async def __anext__(self) -> _DocumentType:
         return await self.next()
 
@@ -1273,7 +1287,11 @@ async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         await self.close()
 
     async def to_list(self) -> list[_DocumentType]:
-        return [x async for x in self]  # noqa: C416,RUF100
+        res: list[_DocumentType] = []
+        while self.alive:
+            if not await self._next_batch(res):
+                break
+        return res
 
 
 class AsyncRawBatchCursor(AsyncCursor, Generic[_DocumentType]):

pymongo/synchronous/command_cursor.py

@@ -346,6 +346,17 @@ def _try_next(self, get_more_allowed: bool) -> Optional[_DocumentType]:
         else:
             return None
 
+    def _next_batch(self, result: list) -> bool:
+        """Get all available documents from the cursor."""
+        if not len(self._data) and not self._killed:
+            self._refresh()
+        if len(self._data):
+            result.extend(self._data)
+            self._data.clear()
+            return True
+        else:
+            return False
+
     def try_next(self) -> Optional[_DocumentType]:
         """Advance the cursor without blocking indefinitely.
 
@@ -371,7 +382,11 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         self.close()
 
     def to_list(self) -> list[_DocumentType]:
-        return [x for x in self]  # noqa: C416,RUF100
+        res: list[_DocumentType] = []
+        while self.alive:
+            if not self._next_batch(res):
+                break
+        return res
 
 
 class RawBatchCommandCursor(CommandCursor[_DocumentType]):

pymongo/synchronous/cursor.py

@@ -1258,6 +1258,20 @@ def next(self) -> _DocumentType:
         else:
             raise StopIteration
 
+    def _next_batch(self, result: list) -> bool:
+        """Get all available documents from the cursor."""
+        if not self._exhaust_checked:
+            self._exhaust_checked = True
+            self._supports_exhaust()
+        if self._empty:
+            return False
+        if len(self._data) or self._refresh():
+            result.extend(self._data)
+            self._data.clear()
+            return True
+        else:
+            return False
+
     def __next__(self) -> _DocumentType:
         return self.next()
 
@@ -1271,7 +1285,11 @@ def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         self.close()
 
     def to_list(self) -> list[_DocumentType]:
-        return [x for x in self]  # noqa: C416,RUF100
+        res: list[_DocumentType] = []
+        while self.alive:
+            if not self._next_batch(res):
+                break
+        return res
 
 
 class RawBatchCursor(Cursor, Generic[_DocumentType]):

test/asynchronous/test_cursor.py

@@ -1380,6 +1380,43 @@ async def test_getMore_does_not_send_readPreference(self):
         self.assertEqual("getMore", started[1].command_name)
         self.assertNotIn("$readPreference", started[1].command)
 
+    @async_client_context.require_replica_set
+    async def test_to_list_tailable(self):
+        oplog = self.client.local.oplog.rs
+        last = await oplog.find().sort("$natural", pymongo.DESCENDING).limit(-1).next()
+        ts = last["ts"]
+
+        c = oplog.find(
+            {"ts": {"$gte": ts}}, cursor_type=pymongo.CursorType.TAILABLE_AWAIT, oplog_replay=True
+        )
+
+        docs = await c.to_list()
+
+        self.assertGreaterEqual(len(docs), 1)
+
+    async def test_to_list_empty(self):
+        c = self.db.does_not_exist.find()
+
+        docs = await c.to_list()
+
+        self.assertEqual([], docs)
+
+    @async_client_context.require_replica_set
+    async def test_command_cursor_to_list(self):
+        c = await self.db.test.aggregate([{"$changeStream": {}}])
+
+        docs = await c.to_list()
+
+        self.assertGreaterEqual(len(docs), 0)
+
+    @async_client_context.require_replica_set
+    async def test_command_cursor_to_list_empty(self):
+        c = await self.db.does_not_exist.aggregate([{"$changeStream": {}}])
+
+        docs = await c.to_list()
+
+        self.assertEqual([], docs)
+
 
 class TestRawBatchCursor(AsyncIntegrationTest):
     async def test_find_raw(self):

test/test_cursor.py

@@ -1371,6 +1371,43 @@ def test_getMore_does_not_send_readPreference(self):
         self.assertEqual("getMore", started[1].command_name)
         self.assertNotIn("$readPreference", started[1].command)
 
+    @client_context.require_replica_set
+    def test_to_list_tailable(self):
+        oplog = self.client.local.oplog.rs
+        last = oplog.find().sort("$natural", pymongo.DESCENDING).limit(-1).next()
+        ts = last["ts"]
+
+        c = oplog.find(
+            {"ts": {"$gte": ts}}, cursor_type=pymongo.CursorType.TAILABLE_AWAIT, oplog_replay=True
+        )
+
+        docs = c.to_list()
+
+        self.assertGreaterEqual(len(docs), 1)
+
+    def test_to_list_empty(self):
+        c = self.db.does_not_exist.find()
+
+        docs = c.to_list()
+
+        self.assertEqual([], docs)
+
+    @client_context.require_replica_set
+    def test_command_cursor_to_list(self):
+        c = self.db.test.aggregate([{"$changeStream": {}}])
+
+        docs = c.to_list()
+
+        self.assertGreaterEqual(len(docs), 0)
+
+    @client_context.require_replica_set
+    def test_command_cursor_to_list_empty(self):
+        c = self.db.does_not_exist.aggregate([{"$changeStream": {}}])
+
+        docs = c.to_list()
+
+        self.assertEqual([], docs)
+
 
 class TestRawBatchCursor(IntegrationTest):
     def test_find_raw(self):