Skip to content

Commit e451ceb

Browse files
committed
fix test errors
1 parent 88ae345 commit e451ceb

File tree

1 file changed

+40
-31
lines changed

1 file changed

+40
-31
lines changed

pymongo/network_layer.py

+40-31
Original file line numberDiff line numberDiff line change
@@ -46,22 +46,21 @@
4646
_HAVE_SSL = False
4747

4848
try:
49-
from pymongo.pyopenssl_context import (
50-
BLOCKING_IO_LOOKUP_ERROR,
51-
BLOCKING_IO_READ_ERROR,
52-
BLOCKING_IO_WRITE_ERROR,
53-
_sslConn,
54-
)
49+
from pymongo.pyopenssl_context import _sslConn as _pysslConn
5550

5651
_HAVE_PYOPENSSL = True
5752
except ImportError:
5853
_HAVE_PYOPENSSL = False
59-
_sslConn = SSLSocket # type: ignore
60-
from pymongo.ssl_support import ( # type: ignore[assignment]
61-
BLOCKING_IO_LOOKUP_ERROR,
62-
BLOCKING_IO_READ_ERROR,
63-
BLOCKING_IO_WRITE_ERROR,
64-
)
54+
_pysslCon = SSLSocket
55+
56+
from pymongo.ssl_support import ( # type: ignore[assignment]
57+
BLOCKING_IO_LOOKUP_ERROR,
58+
BLOCKING_IO_READ_ERROR,
59+
BLOCKING_IO_WRITE_ERROR,
60+
PYBLOCKING_IO_LOOKUP_ERROR,
61+
PYBLOCKING_IO_READ_ERROR,
62+
PYBLOCKING_IO_WRITE_ERROR,
63+
)
6564

6665
if TYPE_CHECKING:
6766
from pymongo.asynchronous.pool import AsyncConnection
@@ -71,17 +70,23 @@
7170
_UNPACK_COMPRESSION_HEADER = struct.Struct("<iiB").unpack
7271
_POLL_TIMEOUT = 0.5
7372
# Errors raised by sockets (and TLS sockets) when in non-blocking mode.
74-
BLOCKING_IO_ERRORS = (BlockingIOError, BLOCKING_IO_LOOKUP_ERROR, *ssl_support.BLOCKING_IO_ERRORS)
73+
BLOCKING_IO_ERRORS = (
74+
BlockingIOError,
75+
BLOCKING_IO_LOOKUP_ERROR,
76+
PYBLOCKING_IO_LOOKUP_ERROR,
77+
*ssl_support.BLOCKING_IO_ERRORS,
78+
*ssl_support.PYBLOCKING_IO_ERRORS,
79+
)
7580

7681

7782
# These socket-based I/O methods are for KMS requests and any other network operations that do not use
7883
# the MongoDB wire protocol
79-
async def async_socket_sendall(sock: Union[socket.socket, _sslConn], buf: bytes) -> None:
84+
async def async_socket_sendall(sock: Union[socket.socket, _pysslConn], buf: bytes) -> None:
8085
timeout = sock.gettimeout()
8186
sock.settimeout(0.0)
8287
loop = asyncio.get_running_loop()
8388
try:
84-
if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
89+
if _HAVE_SSL and isinstance(sock, (SSLSocket, _pysslConn)):
8590
await asyncio.wait_for(_async_socket_sendall_ssl(sock, buf, loop), timeout=timeout)
8691
else:
8792
await asyncio.wait_for(loop.sock_sendall(sock, buf), timeout=timeout) # type: ignore[arg-type]
@@ -95,7 +100,7 @@ async def async_socket_sendall(sock: Union[socket.socket, _sslConn], buf: bytes)
95100
if sys.platform != "win32":
96101

97102
async def _async_socket_sendall_ssl(
98-
sock: Union[socket.socket, _sslConn], buf: bytes, loop: AbstractEventLoop
103+
sock: Union[socket.socket, _pysslConn], buf: bytes, loop: AbstractEventLoop
99104
) -> None:
100105
view = memoryview(buf)
101106
sent = 0
@@ -113,21 +118,23 @@ def _is_ready(fut: Future) -> None:
113118
# Check for closed socket.
114119
if fd == -1:
115120
raise SSLError("Underlying socket has been closed") from None
116-
if isinstance(exc, BLOCKING_IO_READ_ERROR):
121+
if isinstance(exc, (BLOCKING_IO_READ_ERROR, PYBLOCKING_IO_READ_ERROR)):
117122
fut = loop.create_future()
118123
loop.add_reader(fd, _is_ready, fut)
119124
try:
120125
await fut
121126
finally:
122127
loop.remove_reader(fd)
123-
if isinstance(exc, BLOCKING_IO_WRITE_ERROR):
128+
if isinstance(exc, (BLOCKING_IO_WRITE_ERROR, PYBLOCKING_IO_WRITE_ERROR)):
124129
fut = loop.create_future()
125130
loop.add_writer(fd, _is_ready, fut)
126131
try:
127132
await fut
128133
finally:
129134
loop.remove_writer(fd)
130-
if _HAVE_PYOPENSSL and isinstance(exc, BLOCKING_IO_LOOKUP_ERROR):
135+
if _HAVE_PYOPENSSL and isinstance(
136+
exc, (BLOCKING_IO_LOOKUP_ERROR, PYBLOCKING_IO_LOOKUP_ERROR)
137+
):
131138
fut = loop.create_future()
132139
loop.add_reader(fd, _is_ready, fut)
133140
try:
@@ -138,7 +145,7 @@ def _is_ready(fut: Future) -> None:
138145
loop.remove_writer(fd)
139146

140147
async def _async_socket_receive_ssl(
141-
conn: _sslConn, length: int, loop: AbstractEventLoop, once: Optional[bool] = False
148+
conn: _pysslConn, length: int, loop: AbstractEventLoop, once: Optional[bool] = False
142149
) -> memoryview:
143150
mv = memoryview(bytearray(length))
144151
total_read = 0
@@ -162,21 +169,23 @@ def _is_ready(fut: Future) -> None:
162169
# Check for closed socket.
163170
if fd == -1:
164171
raise SSLError("Underlying socket has been closed") from None
165-
if isinstance(exc, BLOCKING_IO_READ_ERROR):
172+
if isinstance(exc, (BLOCKING_IO_READ_ERROR, PYBLOCKING_IO_READ_ERROR)):
166173
fut = loop.create_future()
167174
loop.add_reader(fd, _is_ready, fut)
168175
try:
169176
await fut
170177
finally:
171178
loop.remove_reader(fd)
172-
if isinstance(exc, BLOCKING_IO_WRITE_ERROR):
179+
if isinstance(exc, (BLOCKING_IO_WRITE_ERROR, PYBLOCKING_IO_WRITE_ERROR)):
173180
fut = loop.create_future()
174181
loop.add_writer(fd, _is_ready, fut)
175182
try:
176183
await fut
177184
finally:
178185
loop.remove_writer(fd)
179-
if _HAVE_PYOPENSSL and isinstance(exc, BLOCKING_IO_LOOKUP_ERROR):
186+
if _HAVE_PYOPENSSL and isinstance(
187+
exc, (BLOCKING_IO_LOOKUP_ERROR, PYBLOCKING_IO_LOOKUP_ERROR)
188+
):
180189
fut = loop.create_future()
181190
loop.add_reader(fd, _is_ready, fut)
182191
try:
@@ -192,7 +201,7 @@ def _is_ready(fut: Future) -> None:
192201
# https://docs.python.org/3/library/asyncio-platforms.html#asyncio-platform-support
193202
# Note: In PYTHON-4493 we plan to replace this code with asyncio streams.
194203
async def _async_socket_sendall_ssl(
195-
sock: Union[socket.socket, _sslConn], buf: bytes, dummy: AbstractEventLoop
204+
sock: Union[socket.socket, _pysslConn], buf: bytes, dummy: AbstractEventLoop
196205
) -> None:
197206
view = memoryview(buf)
198207
total_length = len(buf)
@@ -213,7 +222,7 @@ async def _async_socket_sendall_ssl(
213222
total_sent += sent
214223

215224
async def _async_socket_receive_ssl(
216-
conn: _sslConn, length: int, dummy: AbstractEventLoop, once: Optional[bool] = False
225+
conn: _pysslConn, length: int, dummy: AbstractEventLoop, once: Optional[bool] = False
217226
) -> memoryview:
218227
mv = memoryview(bytearray(length))
219228
total_read = 0
@@ -239,7 +248,7 @@ async def _async_socket_receive_ssl(
239248
return mv
240249

241250

242-
def sendall(sock: Union[socket.socket, _sslConn], buf: bytes) -> None:
251+
def sendall(sock: Union[socket.socket, _pysslConn], buf: bytes) -> None:
243252
sock.sendall(buf)
244253

245254

@@ -252,15 +261,15 @@ async def _poll_cancellation(conn: AsyncConnection) -> None:
252261

253262

254263
async def async_receive_data_socket(
255-
sock: Union[socket.socket, _sslConn], length: int
264+
sock: Union[socket.socket, _pysslConn], length: int
256265
) -> memoryview:
257266
sock_timeout = sock.gettimeout()
258267
timeout = sock_timeout
259268

260269
sock.settimeout(0.0)
261270
loop = asyncio.get_running_loop()
262271
try:
263-
if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
272+
if _HAVE_SSL and isinstance(sock, (SSLSocket, _pysslConn)):
264273
return await asyncio.wait_for(
265274
_async_socket_receive_ssl(sock, length, loop, once=True), # type: ignore[arg-type]
266275
timeout=timeout,
@@ -435,7 +444,7 @@ def sock(self) -> socket.socket:
435444

436445

437446
class NetworkingInterface(NetworkingInterfaceBase):
438-
def __init__(self, conn: Union[socket.socket, _sslConn]):
447+
def __init__(self, conn: Union[socket.socket, _pysslConn]):
439448
super().__init__(conn)
440449

441450
def gettimeout(self) -> float | None:
@@ -451,11 +460,11 @@ def is_closing(self) -> bool:
451460
return self.conn.is_closing()
452461

453462
@property
454-
def get_conn(self) -> Union[socket.socket, _sslConn]:
463+
def get_conn(self) -> Union[socket.socket, _pysslConn]:
455464
return self.conn
456465

457466
@property
458-
def sock(self) -> Union[socket.socket, _sslConn]:
467+
def sock(self) -> Union[socket.socket, _pysslConn]:
459468
return self.conn
460469

461470
def fileno(self) -> int:

0 commit comments

Comments
 (0)