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):