Preserve Task types in as_completed

Author: jw9829

stdlib/@tests/test_cases/asyncio/check_as_completed.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import asyncio
+import sys
+from collections.abc import Coroutine
+from typing import Any
+from typing_extensions import assert_type
+
+
+async def result() -> int:
+    return 42
+
+
+if sys.version_info >= (3, 13):
+
+    async def check_as_completed() -> None:
+        task = asyncio.create_task(result())
+
+        async for completed_task in asyncio.as_completed([task]):
+            assert_type(completed_task, asyncio.Task[int])
+            completed_task.cancelling()
+
+        coroutine_for_async_iteration = result()
+        async for completed_future in asyncio.as_completed([coroutine_for_async_iteration]):
+            assert_type(completed_future, asyncio.Future[int])
+
+        for completed_coroutine in asyncio.as_completed([task]):
+            coroutine: Coroutine[Any, Any, int] = completed_coroutine
+            assert_type(await coroutine, int)

stdlib/asyncio/tasks.pyi

@@ -67,6 +67,8 @@ else:
 
 _T = TypeVar("_T")
 _T_co = TypeVar("_T_co", covariant=True)
+_T_sync_co = TypeVar("_T_sync_co", covariant=True)
+_T_async_co = TypeVar("_T_async_co", covariant=True)
 _T1 = TypeVar("_T1")
 _T2 = TypeVar("_T2")
 _T3 = TypeVar("_T3")
@@ -87,9 +89,16 @@ ALL_COMPLETED: Final = concurrent.futures.ALL_COMPLETED
 
 if sys.version_info >= (3, 13):
     @type_check_only
-    class _SyncAndAsyncIterator(Iterator[Coroutine[Any, Any, _T]], AsyncIterator[Future[_T]], Protocol[_T]): ...
+    class _SyncAndAsyncIterator(Iterator[_T_sync_co], AsyncIterator[_T_async_co], Protocol[_T_sync_co, _T_async_co]): ...
 
-    def as_completed(fs: Iterable[_FutureLike[_T]], *, timeout: float | None = None) -> _SyncAndAsyncIterator[_T]: ...
+    @overload
+    def as_completed(
+        fs: Iterable[Task[_T]], *, timeout: float | None = None
+    ) -> _SyncAndAsyncIterator[Coroutine[Any, Any, _T], Task[_T]]: ...
+    @overload
+    def as_completed(
+        fs: Iterable[_FutureLike[_T]], *, timeout: float | None = None
+    ) -> _SyncAndAsyncIterator[Coroutine[Any, Any, _T], Future[_T]]: ...
 
 else:
     def as_completed(fs: Iterable[_FutureLike[_T]], *, timeout: float | None = None) -> Iterator[Future[_T]]: ...