Fix mypy

aio-libs · Oct 28, 2024 · 6e35481 · 6e35481
1 parent e19e929
commit 6e35481
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 4 deletions.
diff --git a/aioloop_proxy/_loop.py b/aioloop_proxy/_loop.py
@@ -17,6 +17,7 @@
     Any,
     ParamSpec,
     Protocol,
+    TypeAlias,
     TypeVar,
     Union,
     cast,
@@ -34,7 +35,14 @@
 _T = TypeVar("_T")
 _Ts = TypeVarTuple("_Ts")
 
-_Coro = Union[Coroutine[Any, Any, _T], Generator[Any, None, _T]]
+_T_co = TypeVar("_T_co", covariant=True)
+
+if sys.version_info >= (3, 12):
+    _AwaitableLike: TypeAlias = Awaitable[_T_co]
+    _CoroutineLike: TypeAlias = Coroutine[Any, Any, _T_co]
+else:
+    _AwaitableLike: TypeAlias = Generator[Any, None, _T_co] | Awaitable[_T_co]
+    _CoroutineLike: TypeAlias = Generator[Any, None, _T_co] | Coroutine[Any, Any, _T_co]
 
 
 # stable
@@ -46,7 +54,7 @@ class _TaskFactory(Protocol):
     def __call__(
         self,
         loop: asyncio.AbstractEventLoop,
-        factory: Coroutine[Any, Any, _T] | Generator[Any, None, _T],
+        factory: _CoroutineLike[_T],
         /,
     ) -> asyncio.Future[_T]: ...
 

diff --git a/aioloop_proxy/_server.py b/aioloop_proxy/_server.py
@@ -2,6 +2,7 @@
 
 import asyncio
 import socket
+import sys
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -37,3 +38,11 @@ async def serve_forever(self) -> None:
 
     async def wait_closed(self) -> None:
         await self._loop._wrap_async(self._orig.wait_closed())
+
+    if sys.version_info >= (3, 13):
+
+        def close_clients(self) -> None:
+            self._loop._wrap_cb(self._orig.close_clients)
+
+        def abort_clients(self) -> None:
+            self._loop._wrap_cb(self._orig.abort_clients)
diff --git a/tests/test_loop.py b/tests/test_loop.py
@@ -11,7 +11,7 @@
 from unittest import mock
 
 import aioloop_proxy
-from aioloop_proxy._loop import _T, _Coro, _ExceptionContext
+from aioloop_proxy._loop import _T, _CoroutineLike, _ExceptionContext
 
 _loop: asyncio.AbstractEventLoop | None = None
 
@@ -110,7 +110,7 @@ def test_task_factory(self) -> None:
 
             def factory(
                 loop: asyncio.AbstractEventLoop,
-                coro: _Coro[_T],
+                coro: _CoroutineLike[_T],
             ) -> asyncio.Task[_T]:
                 nonlocal called
                 called = True