From 0499235facfd0c6d1b1a81e546751ea3888f4441 Mon Sep 17 00:00:00 2001
From: grey-cat-1908 <grey.cat1908@gmail.com>
Date: Sat, 12 Mar 2022 16:26:38 +0300
Subject: [PATCH] shard disconnect, reconnect

---
 melisa/core/gateway.py     |  3 +++
 melisa/core/http.py        |  2 --
 melisa/models/app/shard.py | 30 ++++++++++++++++++++++++------
 3 files changed, 27 insertions(+), 8 deletions(-)

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()