Remove async_timeout requirement on python 3.11+ (#570)

This commit is contained in:
J. Nick Koston 2023-10-12 10:03:30 -10:00 committed by GitHub
parent abc9bb04e4
commit f9fa36a89c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 12 additions and 5 deletions

View File

@ -5,13 +5,13 @@ import contextvars
import enum import enum
import logging import logging
import socket import socket
import sys
import time import time
from collections.abc import Coroutine, Iterable from collections.abc import Coroutine, Iterable
from dataclasses import astuple, dataclass from dataclasses import astuple, dataclass
from functools import partial from functools import partial
from typing import TYPE_CHECKING, Any, Callable from typing import TYPE_CHECKING, Any, Callable
import async_timeout
from google.protobuf import message from google.protobuf import message
import aioesphomeapi.host_resolver as hr import aioesphomeapi.host_resolver as hr
@ -45,6 +45,13 @@ from .core import (
) )
from .model import APIVersion from .model import APIVersion
# mypy: disable-error-code="import-not-found"
if sys.version_info[:2] < (3, 11):
from async_timeout import timeout as asyncio_timeout
else:
from asyncio import timeout as asyncio_timeout
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
BUFFER_SIZE = 1024 * 1024 * 2 # Set buffer limit to 2MB BUFFER_SIZE = 1024 * 1024 * 2 # Set buffer limit to 2MB
@ -265,7 +272,7 @@ class APIConnection:
self._params.port, self._params.port,
self._params.zeroconf_instance, self._params.zeroconf_instance,
) )
async with async_timeout.timeout(RESOLVE_TIMEOUT): async with asyncio_timeout(RESOLVE_TIMEOUT):
return await coro return await coro
except asyncio.TimeoutError as err: except asyncio.TimeoutError as err:
raise ResolveAPIError( raise ResolveAPIError(
@ -302,7 +309,7 @@ class APIConnection:
try: try:
coro = self._loop.sock_connect(self._socket, sockaddr) coro = self._loop.sock_connect(self._socket, sockaddr)
async with async_timeout.timeout(TCP_CONNECT_TIMEOUT): async with asyncio_timeout(TCP_CONNECT_TIMEOUT):
await coro await coro
except asyncio.TimeoutError as err: except asyncio.TimeoutError as err:
raise SocketAPIError(f"Timeout while connecting to {sockaddr}") from err raise SocketAPIError(f"Timeout while connecting to {sockaddr}") from err
@ -486,7 +493,7 @@ class APIConnection:
# Allow 2 minutes for connect and setup; this is only as a last measure # Allow 2 minutes for connect and setup; this is only as a last measure
# to protect from issues if some part of the connect process mistakenly # to protect from issues if some part of the connect process mistakenly
# does not have a timeout # does not have a timeout
async with async_timeout.timeout(CONNECT_AND_SETUP_TIMEOUT): async with asyncio_timeout(CONNECT_AND_SETUP_TIMEOUT):
await self._connect_task await self._connect_task
except asyncio.CancelledError: except asyncio.CancelledError:
# If the task was cancelled, we need to clean up the connection # If the task was cancelled, we need to clean up the connection

View File

@ -2,4 +2,4 @@ protobuf>=3.19.0
zeroconf>=0.36.0,<1.0 zeroconf>=0.36.0,<1.0
chacha20poly1305-reuseable>=0.2.5 chacha20poly1305-reuseable>=0.2.5
noiseprotocol>=0.3.1,<1.0 noiseprotocol>=0.3.1,<1.0
async-timeout>=4.0 async-timeout>=4.0;python_version<'3.11'