From ef32f6d8822eb17c41a636094cf4a1232eab966f Mon Sep 17 00:00:00 2001 From: Stocker <44980366+StockerMC@users.noreply.github.com> Date: Fri, 20 Aug 2021 19:50:39 -0400 Subject: [PATCH] Typehint context_managers.py --- discord/context_managers.py | 38 +++++++++++++++++++++++++++---------- 1 file changed, 28 insertions(+), 10 deletions(-) diff --git a/discord/context_managers.py b/discord/context_managers.py index bb3b77ab7..a3ab0d197 100644 --- a/discord/context_managers.py +++ b/discord/context_managers.py @@ -22,13 +22,23 @@ FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. """ +from __future__ import annotations + import asyncio +from typing import TYPE_CHECKING, TypeVar, Optional, Type + +if TYPE_CHECKING: + from .abc import Messageable + + from types import TracebackType + + TypingT = TypeVar('TypingT', bound='Typing') __all__ = ( 'Typing', ) -def _typing_done_callback(fut): +def _typing_done_callback(fut: asyncio.Future) -> None: # just retrieve any exception and call it a day try: fut.exception() @@ -36,11 +46,11 @@ def _typing_done_callback(fut): pass class Typing: - def __init__(self, messageable): - self.loop = messageable._state.loop - self.messageable = messageable + def __init__(self, messageable: Messageable) -> None: + self.loop: asyncio.AbstractEventLoop = messageable._state.loop + self.messageable: Messageable = messageable - async def do_typing(self): + async def do_typing(self) -> None: try: channel = self._channel except AttributeError: @@ -52,18 +62,26 @@ class Typing: await typing(channel.id) await asyncio.sleep(5) - def __enter__(self): - self.task = asyncio.ensure_future(self.do_typing(), loop=self.loop) + def __enter__(self: TypingT) -> TypingT: + self.task: asyncio.Task = self.loop.create_task(self.do_typing()) self.task.add_done_callback(_typing_done_callback) return self - def __exit__(self, exc_type, exc, tb): + def __exit__(self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: self.task.cancel() - async def __aenter__(self): + async def __aenter__(self: TypingT) -> TypingT: self._channel = channel = await self.messageable._get_channel() await channel._state.http.send_typing(channel.id) return self.__enter__() - async def __aexit__(self, exc_type, exc, tb): + async def __aexit__(self, + exc_type: Optional[Type[BaseException]], + exc_value: Optional[BaseException], + traceback: Optional[TracebackType], + ) -> None: self.task.cancel()