Type error fixes

Author: NoahStapp

pymongo/asynchronous/encryption.py

@@ -143,9 +143,14 @@ def __init__(
         self._spawned = False
 
     async def _async_kms_request(
-        self, kms_context: MongoCryptKmsContext, host, port, opts, message
+        self,
+        kms_context: MongoCryptKmsContext,
+        host: str,
+        port: Optional[int],
+        opts: PoolOptions,
+        message: bytes,
     ) -> None:
-        from pymongo.network_layer import _async_receive_data_socket
+        from pymongo.network_layer import async_receive_data_socket  # type: ignore[attr-defined]
 
         try:
             conn = await _configured_socket((host, port), opts)
@@ -154,7 +159,7 @@ async def _async_kms_request(
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
                     conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    data = await _async_receive_data_socket(conn, kms_context.bytes_needed)
+                    data = await async_receive_data_socket(conn, kms_context.bytes_needed)
                     kms_context.feed(data)
             except OSError as err:
                 raise OSError("KMS connection closed") from err

pymongo/network_layer.py

@@ -277,7 +277,9 @@ async def async_receive_data(
         sock.settimeout(sock_timeout)
 
 
-async def _async_receive_data_socket(sock: socket.socket | _sslConn, length: int) -> memoryview:
+async def async_receive_data_socket(
+    sock: Union[socket.socket, _sslConn], length: int
+) -> memoryview:
     sock_timeout = sock.gettimeout()
     timeout = sock_timeout
 
@@ -286,8 +288,9 @@ async def _async_receive_data_socket(sock: socket.socket | _sslConn, length: int
     try:
         if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
             return await asyncio.wait_for(
-                _async_receive_ssl(sock, length, loop, once=True), timeout=timeout
-            )  # type: ignore[arg-type]
+                _async_receive_ssl(sock, length, loop, once=True),  # type: ignore[arg-type]
+                timeout=timeout,
+            )
         else:
             return await asyncio.wait_for(_async_receive(sock, length, loop), timeout=timeout)  # type: ignore[arg-type]
     except asyncio.TimeoutError as err:

pymongo/synchronous/encryption.py

@@ -143,9 +143,14 @@ def __init__(
         self._spawned = False
 
     def _async_kms_request(
-        self, kms_context: MongoCryptKmsContext, host, port, opts, message
+        self,
+        kms_context: MongoCryptKmsContext,
+        host: str,
+        port: Optional[int],
+        opts: PoolOptions,
+        message: bytes,
     ) -> None:
-        from pymongo.network_layer import _receive_data_socket
+        from pymongo.network_layer import receive_data_socket  # type: ignore[attr-defined]
 
         try:
             conn = _configured_socket((host, port), opts)
@@ -154,7 +159,7 @@ def _async_kms_request(
                 while kms_context.bytes_needed > 0:
                     # CSOT: update timeout.
                     conn.settimeout(max(_csot.clamp_remaining(_KMS_CONNECT_TIMEOUT), 0))
-                    data = _receive_data_socket(conn, kms_context.bytes_needed)
+                    data = receive_data_socket(conn, kms_context.bytes_needed)
                     kms_context.feed(data)
             except OSError as err:
                 raise OSError("KMS connection closed") from err