diff --git a/aiojobs/_job.py b/aiojobs/_job.py index ffcf659a..e8e21028 100644 --- a/aiojobs/_job.py +++ b/aiojobs/_job.py @@ -3,7 +3,10 @@ import traceback from typing import TYPE_CHECKING, Coroutine, Generic, Optional, TypeVar -import async_timeout +if sys.version_info >= (3, 11): + from asyncio import timeout as asyncio_timeout +else: + from async_timeout import timeout as asyncio_timeout if TYPE_CHECKING: from ._scheduler import Scheduler @@ -70,7 +73,7 @@ def set_name(self, name: str) -> None: self._task.set_name(name) async def _do_wait(self, timeout: Optional[float]) -> _T: - async with async_timeout.timeout(timeout): + async with asyncio_timeout(timeout): # TODO: add a test for waiting for a pending coro await self._started assert self._task is not None # Task should have been created before this. @@ -113,7 +116,7 @@ async def _close(self, timeout: Optional[float]) -> None: # self._scheduler is None after _done_callback() scheduler = self._scheduler try: - async with async_timeout.timeout(timeout): + async with asyncio_timeout(timeout): await self._task except asyncio.CancelledError: pass diff --git a/setup.cfg b/setup.cfg index 883c033f..bae945ab 100644 --- a/setup.cfg +++ b/setup.cfg @@ -46,7 +46,7 @@ packages = aiojobs include_package_data = True install_requires = - async-timeout >= 4.0.0 + async-timeout >= 4.0.0 ; python_version < "3.11" [options.extras_require] aiohttp = diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py index 7fe760b7..e105feb7 100644 --- a/tests/test_scheduler.py +++ b/tests/test_scheduler.py @@ -1,12 +1,17 @@ import asyncio +import sys from typing import Awaitable, Callable, List, NoReturn from unittest import mock import pytest -from async_timeout import timeout from aiojobs import Scheduler +if sys.version_info >= (3, 11): + from asyncio import timeout as asyncio_timeout +else: + from async_timeout import timeout as asyncio_timeout + _MakeScheduler = Callable[..., Awaitable[Scheduler]] @@ -285,7 +290,7 @@ async def coro(fut: "asyncio.Future[None]") -> None: with pytest.raises(asyncio.TimeoutError): # try to wait for 1 sec to add task to pending queue - async with timeout(1): + async with asyncio_timeout(1): await scheduler.spawn(coro(fut3)) assert scheduler.active_count == 1