Reconnect when WebSocket dies

Check if the WebSocket connection was supposed to terminate and
reconnect to the gateway if this is not the case.

1 files changed, 15 insertions, 0 deletions

diff --git a/discord/client.py b/discord/client.py
index cd5293f4..67fba4ff 100644
--- a/discord/client.py
+++ b/discord/client.py
@@ -352,6 +352,8 @@ class Client(object):
 
     def __init__(self, **kwargs):
         self._is_logged_in = False
+        self._close = False
+        self.options = kwargs
         self.connection = ConnectionState(self.dispatch, **kwargs)
         self.dispatch_lock = threading.RLock()
         self.token = ''
@@ -467,6 +469,18 @@ class Client(object):
         log.info('Client is being run')
         self.ws.run()
 
+        # The WebSocket is guaranteed to be terminated after ws.run().
+        # Check if we wanted it to close and reconnect if not.
+        while not self._close:
+            gateway = requests.get(endpoints.GATEWAY, headers=self.headers)
+            if gateway.status_code != 200:
+                raise GatewayNotFound()
+            self.connection = ConnectionState(self.dispatch, **self.options)
+            self._create_websocket(gateway.json().get('url'), reconnect=False)
+            self.ws.run()
+
+        log.info('Client exiting')
+
     @property
     def is_logged_in(self):
         """Returns True if the client is successfully logged in. False otherwise."""
@@ -642,6 +656,7 @@ class Client(object):
     def logout(self):
         """Logs out of Discord and closes all connections."""
         response = requests.post(endpoints.LOGOUT)
+        self._close = True
         self.ws.close()
         self._is_logged_in = False
         log.debug(request_logging_format.format(name='logout', response=response))