Typing fixes + comment

Author: NoahStapp

pymongo/asynchronous/encryption.py

@@ -183,8 +183,8 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     if _IS_SYNC:
                         data = conn.recv(kms_context.bytes_needed)
                     else:
-                        from pymongo.network_layer import (
-                            async_receive_data_socket,  # type: ignore[attr-defined]
+                        from pymongo.network_layer import (  # type: ignore[attr-defined]
+                            async_receive_data_socket,
                         )
 
                         data = await async_receive_data_socket(conn, kms_context.bytes_needed)

pymongo/network_layer.py

@@ -145,6 +145,7 @@ def _is_ready(fut: Future) -> None:
                 read = conn.recv_into(mv[total_read:])
                 if read == 0:
                     raise OSError("connection closed")
+                # KMS responses update their expected size after the first batch, stop reading after one loop
                 if once:
                     return mv[:read]
                 total_read += read

pymongo/synchronous/encryption.py

@@ -183,8 +183,8 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                     if _IS_SYNC:
                         data = conn.recv(kms_context.bytes_needed)
                     else:
-                        from pymongo.network_layer import (
-                            receive_data_socket,  # type: ignore[attr-defined]
+                        from pymongo.network_layer import (  # type: ignore[attr-defined]
+                            receive_data_socket,
                         )
 
                         data = receive_data_socket(conn, kms_context.bytes_needed)