fix different socket types on 3.6/7 vs >=3.8

This commit is contained in:
Jack Robison 2020-10-30 12:59:48 -04:00
parent a33a3a3a83
commit 0e90927f0e
No known key found for this signature in database
GPG key ID: DF25C68FE0239BB2

View file

@ -1,17 +1,24 @@
import sys
import struct
import socket
import typing
from asyncio.protocols import DatagramProtocol
from asyncio.transports import DatagramTransport
from asyncio.trsock import TransportSocket
from unittest import mock
if sys.version_info >= (3, 8):
from asyncio.trsock import TransportSocket
SOCKET_TYPES = (socket.SocketType, TransportSocket, mock.MagicMock)
else:
SOCKET_TYPES = (socket.SocketType, mock.MagicMock)
def _get_sock(transport: typing.Optional[DatagramTransport]) -> typing.Optional[socket.socket]:
if transport is None or not hasattr(transport, "_extra"):
return None
sock: typing.Optional[socket.socket] = transport.get_extra_info('socket', None)
assert sock is None or isinstance(sock, (socket.SocketType, TransportSocket, mock.MagicMock))
assert sock is None or isinstance(sock, SOCKET_TYPES)
return sock