1414)
1515from sys import platform , version_info
1616from time import perf_counter
17- from typing import Any , Dict , List , Optional , Tuple , Union
17+ from typing import TYPE_CHECKING , Any , Dict , List , Optional , Tuple , Type , Union
1818
19- from aiohttp import WSMessage , WSMsgType
19+ from aiohttp import ClientWebSocketResponse , WSMessage , WSMsgType
2020from aiohttp .http import WS_CLOSED_MESSAGE , WS_CLOSING_MESSAGE
2121
2222from ...base import get_logger
3131from ..models .presence import ClientPresence
3232from .heartbeat import _Heartbeat
3333
34+ if TYPE_CHECKING :
35+ from ...client .context import _Context
36+
3437log = get_logger ("gateway" )
3538
3639__all__ = ("WebSocketClient" ,)
@@ -102,30 +105,30 @@ def __init__(
102105 self ._loop = get_event_loop () if version_info < (3 , 10 ) else get_running_loop ()
103106 except RuntimeError :
104107 self ._loop = new_event_loop ()
105- self ._dispatch = Listener ()
106- self ._http = HTTPClient (token )
107- self ._client = None
108- self ._closed = False
109- self ._options = {
108+ self ._dispatch : Listener = Listener ()
109+ self ._http : HTTPClient = HTTPClient (token )
110+ self ._client : Optional [ "ClientWebSocketResponse" ] = None
111+ self ._closed : bool = False
112+ self ._options : dict = {
110113 "max_msg_size" : 1024 ** 2 ,
111114 "timeout" : 60 ,
112115 "autoclose" : False ,
113116 "compress" : 0 ,
114117 }
115- self ._intents = intents
118+ self ._intents : Intents = intents
116119 self .__heartbeater : _Heartbeat = _Heartbeat (
117120 loop = self ._loop if version_info < (3 , 10 ) else None
118121 )
119- self .__shard = None
120- self .__presence = None
121- self .__task = None
122- self .session_id = None if session_id is MISSING else session_id
123- self .sequence = None if sequence is MISSING else sequence
124- self .ready = Event (loop = self ._loop ) if version_info < (3 , 10 ) else Event ()
125-
126- self ._last_send = perf_counter ()
127- self ._last_ack = perf_counter ()
128- self .latency : float ("nan" ) # noqa: F821
122+ self .__shard : Optional [ List [ Tuple [ int ]]] = None
123+ self .__presence : Optional [ ClientPresence ] = None
124+ self .__task : Optional [ Task ] = None
125+ self .session_id : Optional [ str ] = None if session_id is MISSING else session_id
126+ self .sequence : Optional [ str ] = None if sequence is MISSING else sequence
127+ self .ready : Event = Event (loop = self ._loop ) if version_info < (3 , 10 ) else Event ()
128+
129+ self ._last_send : float = perf_counter ()
130+ self ._last_ack : float = perf_counter ()
131+ self .latency : float = float ("nan" ) # noqa: F821
129132 # self.latency has to be noqa, this is valid in python but not in Flake8.
130133
131134 async def _manage_heartbeat (self ) -> None :
@@ -142,10 +145,9 @@ async def _manage_heartbeat(self) -> None:
142145 await self .__restart ()
143146 break
144147
145- async def __restart (self ):
148+ async def __restart (self ) -> None :
146149 """Restart the client's connection and heartbeat with the Gateway."""
147150 if self .__task :
148- self .__task : Task
149151 self .__task .cancel ()
150152 self ._client = None # clear pending waits
151153 self .__heartbeater .event .clear ()
@@ -258,7 +260,7 @@ async def _handle_connection(
258260 log .debug (f"{ event } : { data } " )
259261 self ._dispatch_event (event , data )
260262
261- async def wait_until_ready (self ):
263+ async def wait_until_ready (self ) -> None :
262264 """Waits for the client to become ready according to the Gateway."""
263265 await self .ready .wait ()
264266
@@ -395,15 +397,15 @@ def _dispatch_event(self, event: str, data: dict) -> None: # sourcery no-metric
395397 log .fatal (f"An error occured dispatching { name } : { error } " )
396398 self ._dispatch .dispatch ("raw_socket_create" , data )
397399
398- def __contextualize (self , data : dict ) -> object :
400+ def __contextualize (self , data : dict ) -> "_Context" :
399401 """
400402 Takes raw data given back from the Gateway
401403 and gives "context" based off of what it is.
402404
403405 :param data: The data from the Gateway.
404406 :type data: dict
405407 :return: The context object.
406- :rtype: object
408+ :rtype: Any
407409 """
408410 if data ["type" ] != InteractionType .PING :
409411 _context : str = ""
@@ -418,12 +420,12 @@ def __contextualize(self, data: dict) -> object:
418420 _context = "ComponentContext"
419421
420422 data ["_client" ] = self ._http
421- context : object = getattr (__import__ ("interactions.client.context" ), _context )
423+ context : Type [ "_Context" ] = getattr (__import__ ("interactions.client.context" ), _context )
422424
423425 return context (** data )
424426
425427 def __sub_command_context (
426- self , data : Union [dict , Option ], context : object
428+ self , data : Union [dict , Option ], context : "_Context"
427429 ) -> Union [Tuple [str ], dict ]:
428430 """
429431 Checks if an application command schema has sub commands
@@ -510,7 +512,7 @@ def _check_auto(option: dict) -> Optional[Tuple[str]]:
510512
511513 return __kwargs
512514
513- def __option_type_context (self , context : object , type : int ) -> dict :
515+ def __option_type_context (self , context : "_Context" , type : int ) -> dict :
514516 """
515517 Looks up the type of option respective to the existing
516518 option types.
@@ -544,11 +546,11 @@ def __option_type_context(self, context: object, type: int) -> dict:
544546 }
545547 return _resolved
546548
547- async def restart (self ):
549+ async def restart (self ) -> None :
548550 await self .__restart ()
549551
550552 @property
551- async def __receive_packet_stream (self ) -> Optional [Dict [str , Any ]]:
553+ async def __receive_packet_stream (self ) -> Optional [Union [ Dict [str , Any ], WSMessage ]]:
552554 """
553555 Receives a stream of packets sent from the Gateway.
554556
0 commit comments