diff --git a/melisa/core/gateway.py b/melisa/core/gateway.py index 3f303ca..94ccdb8 100644 --- a/melisa/core/gateway.py +++ b/melisa/core/gateway.py @@ -47,6 +47,7 @@ class Gateway: self.client = client self.shard_id = shard_id self.latency = float('inf') + self.connected = False self.listeners = listeners @@ -87,6 +88,7 @@ class Gateway: await self.hello() if self.interval is None: return + self.connected = True await asyncio.gather(self.heartbeat(), self.receive()) async def close(self, code: int = 1000): @@ -136,6 +138,7 @@ class Gateway: async def heartbeat(self): while self.interval is not None: await self.send(self.HEARTBEAT, self.sequence) + self.connected = True await asyncio.sleep(self.interval) async def hello(self): diff --git a/melisa/core/http.py b/melisa/core/http.py index b009011..e668512 100644 --- a/melisa/core/http.py +++ b/melisa/core/http.py @@ -143,5 +143,3 @@ class HTTPClient: route, headers=headers ) - - diff --git a/melisa/models/app/shard.py b/melisa/models/app/shard.py index a33cdc1..a725390 100644 --- a/melisa/models/app/shard.py +++ b/melisa/models/app/shard.py @@ -1,6 +1,7 @@ from __future__ import annotations -from asyncio import create_task +from asyncio import create_task, Task, sleep +from typing import Optional from ...core.gateway import Gateway from ..user import BotActivity @@ -24,8 +25,11 @@ class Shard: """Id of Shard""" return self._gateway.shard_id + # @property + # def + async def launch(self, **kwargs) -> Shard: - """Launch new shard""" + """Launches new shard""" self._gateway = Gateway(self._client, self._shard_id, self._num_shards, @@ -40,6 +44,11 @@ class Shard: return self + async def _try_close(self) -> None: + if self._gateway.connected: + self._gateway.connected = False + await self._gateway.close(code=1000) + async def update_presence(self, activity: BotActivity = None, status: str = None) -> Shard: """ Update Presence for the shard @@ -56,9 +65,18 @@ class Shard: return self - async def disconnect(self) -> Shard: - await self._gateway.close() + async def disconnect(self) -> None: + """Disconnect current shard""" + await self._try_close() - self.disconnected = True + async def reconnect(self, wait_time: int = 3) -> None: + """Reconnect current shard - return self + Parameters + ---------- + wait_time: :class:`int` + Reconnect after + """ + await self._try_close() + await sleep(wait_time) + await self.launch()