From 190e07236f38298e281e7810d304a1313f91668a Mon Sep 17 00:00:00 2001
From: Sebastian Rittau <srittau@rittau.biz>
Date: Thu, 8 Apr 2021 17:31:45 +0200
Subject: [PATCH] Tighten ensure_future() return type using overloads (#5190)

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
---
 stdlib/asyncio/tasks.pyi | 24 +++++++-----------------
 1 file changed, 7 insertions(+), 17 deletions(-)

diff --git a/stdlib/asyncio/tasks.pyi b/stdlib/asyncio/tasks.pyi
index f8070017c..9da1d92f5 100644
--- a/stdlib/asyncio/tasks.pyi
+++ b/stdlib/asyncio/tasks.pyi
@@ -1,22 +1,8 @@
 import concurrent.futures
 import sys
+from collections.abc import Awaitable, Generator, Iterable, Iterator
 from types import FrameType
-from typing import (
-    Any,
-    Awaitable,
-    Generator,
-    Generic,
-    Iterable,
-    Iterator,
-    List,
-    Optional,
-    Set,
-    TextIO,
-    Tuple,
-    TypeVar,
-    Union,
-    overload,
-)
+from typing import Any, Generic, List, Optional, Set, TextIO, Tuple, TypeVar, Union, overload
 from typing_extensions import Literal
 
 from .events import AbstractEventLoop
@@ -31,6 +17,7 @@ _T2 = TypeVar("_T2")
 _T3 = TypeVar("_T3")
 _T4 = TypeVar("_T4")
 _T5 = TypeVar("_T5")
+_FT = TypeVar("_FT", bound=Future[Any])
 _FutureT = Union[Future[_T], Generator[Any, None, _T], Awaitable[_T]]
 _TaskYieldType = Optional[Future[object]]
 
@@ -41,7 +28,10 @@ ALL_COMPLETED: str
 def as_completed(
     fs: Iterable[_FutureT[_T]], *, loop: Optional[AbstractEventLoop] = ..., timeout: Optional[float] = ...
 ) -> Iterator[Future[_T]]: ...
-def ensure_future(coro_or_future: _FutureT[_T], *, loop: Optional[AbstractEventLoop] = ...) -> Future[_T]: ...
+@overload
+def ensure_future(coro_or_future: _FT, *, loop: Optional[AbstractEventLoop] = ...) -> _FT: ...  # type: ignore
+@overload
+def ensure_future(coro_or_future: Awaitable[_T], *, loop: Optional[AbstractEventLoop] = ...) -> Task[_T]: ...
 
 # Prior to Python 3.7 'async' was an alias for 'ensure_future'.
 # It became a keyword in 3.7.