Merge branch 'master' of github.com:mongodb/mongo-python-driver

Author: blink1073

doc/api/pymongo/asynchronous/mongo_client.rst

@@ -6,7 +6,7 @@
 
    .. autoclass:: pymongo.asynchronous.mongo_client.AsyncMongoClient(host='localhost', port=27017, document_class=dict, tz_aware=False, connect=True, **kwargs)
 
-      .. automethod:: aclose
+      .. automethod:: close
 
       .. describe:: c[db_name] || c.db_name
 

gridfs/asynchronous/grid_file.py

@@ -1892,8 +1892,16 @@ 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
+    async def to_list(self, length: Optional[int] = None) -> list[AsyncGridOut]:
+        """Convert the cursor to a list."""
+        if length is None:
+            return [x async for x in self]  # noqa: C416,RUF100
+        if length < 1:
+            raise ValueError("to_list() length must be greater than 0")
+        ret = []
+        for _ in range(length):
+            ret.append(await self.next())
+        return ret
 
     __anext__ = next
 

gridfs/synchronous/grid_file.py

@@ -1878,8 +1878,16 @@ 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
+    def to_list(self, length: Optional[int] = None) -> list[GridOut]:
+        """Convert the cursor to a list."""
+        if length is None:
+            return [x for x in self]  # noqa: C416,RUF100
+        if length < 1:
+            raise ValueError("to_list() length must be greater than 0")
+        ret = []
+        for _ in range(length):
+            ret.append(self.next())
+        return ret
 
     __next__ = next
 

pymongo/asynchronous/collection.py

@@ -1893,9 +1893,7 @@ def find(self, *args: Any, **kwargs: Any) -> AsyncCursor[_DocumentType]:
         """
         return AsyncCursor(self, *args, **kwargs)
 
-    async def find_raw_batches(
-        self, *args: Any, **kwargs: Any
-    ) -> AsyncRawBatchCursor[_DocumentType]:
+    def find_raw_batches(self, *args: Any, **kwargs: Any) -> AsyncRawBatchCursor[_DocumentType]:
         """Query the database and retrieve batches of raw BSON.
 
         Similar to the :meth:`find` method but returns a
@@ -1907,7 +1905,7 @@ async def find_raw_batches(
         :mod:`bson` module.
 
           >>> import bson
-          >>> cursor = await db.test.find_raw_batches()
+          >>> cursor = db.test.find_raw_batches()
           >>> async for batch in cursor:
           ...     print(bson.decode_all(batch))
 

pymongo/asynchronous/command_cursor.py

@@ -346,13 +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."""
+    async def _next_batch(self, result: list, total: Optional[int] = None) -> bool:
+        """Get all or some 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()
+            if total is None:
+                result.extend(self._data)
+                self._data.clear()
+            else:
+                for _ in range(min(len(self._data), total)):
+                    result.append(self._data.popleft())
             return True
         else:
             return False
@@ -381,21 +385,32 @@ async def __aenter__(self) -> AsyncCommandCursor[_DocumentType]:
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         await self.close()
 
-    async def to_list(self) -> list[_DocumentType]:
+    async def to_list(self, length: Optional[int] = None) -> list[_DocumentType]:
         """Converts the contents of this cursor to a list more efficiently than ``[doc async for doc in cursor]``.
 
         To use::
 
           >>> await cursor.to_list()
 
+        Or, so read at most n items from the cursor::
+
+          >>> await cursor.to_list(n)
+
         If the cursor is empty or has no more results, an empty list will be returned.
 
         .. versionadded:: 4.9
         """
         res: list[_DocumentType] = []
+        remaining = length
+        if isinstance(length, int) and length < 1:
+            raise ValueError("to_list() length must be greater than 0")
         while self.alive:
-            if not await self._next_batch(res):
+            if not await self._next_batch(res, remaining):
                 break
+            if length is not None:
+                remaining = length - len(res)
+                if remaining == 0:
+                    break
         return res
 
 

pymongo/asynchronous/cursor.py

@@ -1260,16 +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."""
+    async def _next_batch(self, result: list, total: Optional[int] = None) -> bool:
+        """Get all or some 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()
+            if total is None:
+                result.extend(self._data)
+                self._data.clear()
+            else:
+                for _ in range(min(len(self._data), total)):
+                    result.append(self._data.popleft())
             return True
         else:
             return False
@@ -1286,21 +1290,32 @@ async def __aenter__(self) -> AsyncCursor[_DocumentType]:
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         await self.close()
 
-    async def to_list(self) -> list[_DocumentType]:
+    async def to_list(self, length: Optional[int] = None) -> list[_DocumentType]:
         """Converts the contents of this cursor to a list more efficiently than ``[doc async for doc in cursor]``.
 
         To use::
 
           >>> await cursor.to_list()
 
+        Or, so read at most n items from the cursor::
+
+          >>> await cursor.to_list(n)
+
         If the cursor is empty or has no more results, an empty list will be returned.
 
         .. versionadded:: 4.9
         """
         res: list[_DocumentType] = []
+        remaining = length
+        if isinstance(length, int) and length < 1:
+            raise ValueError("to_list() length must be greater than 0")
         while self.alive:
-            if not await self._next_batch(res):
+            if not await self._next_batch(res, remaining):
                 break
+            if length is not None:
+                remaining = length - len(res)
+                if remaining == 0:
+                    break
         return res
 
 

pymongo/asynchronous/encryption.py

@@ -299,7 +299,7 @@ async def close(self) -> None:
         self.client_ref = None
         self.key_vault_coll = None
         if self.mongocryptd_client:
-            await self.mongocryptd_client.aclose()
+            await self.mongocryptd_client.close()
             self.mongocryptd_client = None
 
 
@@ -439,7 +439,7 @@ async def close(self) -> None:
         self._closed = True
         await self._auto_encrypter.close()
         if self._internal_client:
-            await self._internal_client.aclose()
+            await self._internal_client.close()
             self._internal_client = None
 
 

pymongo/asynchronous/mongo_client.py

@@ -1378,7 +1378,7 @@ async def __aenter__(self) -> AsyncMongoClient[_DocumentType]:
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        await self.aclose()
+        await self.close()
 
     # See PYTHON-3084.
     __iter__ = None
@@ -1514,7 +1514,7 @@ async def _end_sessions(self, session_ids: list[_ServerSession]) -> None:
             # command.
             pass
 
-    async def aclose(self) -> None:
+    async def close(self) -> None:
         """Cleanup client resources and disconnect from MongoDB.
 
         End all server sessions created by this client by sending one or more
@@ -1541,6 +1541,10 @@ async def aclose(self) -> None:
             # TODO: PYTHON-1921 Encrypted MongoClients cannot be re-opened.
             await self._encrypter.close()
 
+    if not _IS_SYNC:
+        # Add support for contextlib.aclosing.
+        aclose = close
+
     async def _get_topology(self) -> Topology:
         """Get the internal :class:`~pymongo.asynchronous.topology.Topology` object.
 

pymongo/network_layer.py

@@ -77,6 +77,7 @@ async def async_sendall(sock: Union[socket.socket, _sslConn], buf: bytes) -> Non
 async def _async_sendall_ssl(
     sock: Union[socket.socket, _sslConn], buf: bytes, loop: AbstractEventLoop
 ) -> None:
+    view = memoryview(buf)
     fd = sock.fileno()
     sent = 0
 
@@ -89,7 +90,7 @@ def _is_ready(fut: Future) -> None:
 
     while sent < len(buf):
         try:
-            sent += sock.send(buf)
+            sent += sock.send(view[sent:])
         except BLOCKING_IO_ERRORS as exc:
             fd = sock.fileno()
             # Check for closed socket.

pymongo/synchronous/command_cursor.py

@@ -346,13 +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."""
+    def _next_batch(self, result: list, total: Optional[int] = None) -> bool:
+        """Get all or some 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()
+            if total is None:
+                result.extend(self._data)
+                self._data.clear()
+            else:
+                for _ in range(min(len(self._data), total)):
+                    result.append(self._data.popleft())
             return True
         else:
             return False
@@ -381,21 +385,32 @@ def __enter__(self) -> CommandCursor[_DocumentType]:
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
         self.close()
 
-    def to_list(self) -> list[_DocumentType]:
+    def to_list(self, length: Optional[int] = None) -> list[_DocumentType]:
         """Converts the contents of this cursor to a list more efficiently than ``[doc for doc in cursor]``.
 
         To use::
 
           >>> cursor.to_list()
 
+        Or, so read at most n items from the cursor::
+
+          >>> cursor.to_list(n)
+
         If the cursor is empty or has no more results, an empty list will be returned.
 
         .. versionadded:: 4.9
         """
         res: list[_DocumentType] = []
+        remaining = length
+        if isinstance(length, int) and length < 1:
+            raise ValueError("to_list() length must be greater than 0")
         while self.alive:
-            if not self._next_batch(res):
+            if not self._next_batch(res, remaining):
                 break
+            if length is not None:
+                remaining = length - len(res)
+                if remaining == 0:
+                    break
         return res