Add .close coroutine to BfxWebsocketClient and _BfxWebsocketBucket classes.

bfxapi/websocket/BfxWebsocketClient.py

@@ -44,11 +44,12 @@ class BfxWebsocketClient(object):
         await asyncio.gather(*tasks)
 
     async def __connect(self, API_KEY, API_SECRET, filter=None):
-        async with websockets.connect(self.host) as websocket:
+        async for websocket in websockets.connect(self.host):
             self.websocket = websocket
             
             await self.__authenticate(API_KEY, API_SECRET, filter)
 
+            try:
                 async for message in websocket:
                     message = json.loads(message)
 
@@ -60,6 +61,7 @@ class BfxWebsocketClient(object):
                         self.event_emitter.emit("wss-error", message["code"], message["msg"])
                     elif isinstance(message, list) and (chanId := message[0]) == 0 and message[1] != HEARTBEAT:
                         self.handler.handle(message[1], message[2])
+            except websockets.ConnectionClosedError: continue
 
     async def __authenticate(self, API_KEY, API_SECRET, filter=None):
         data = { "event": "auth", "filter": filter, "apiKey": API_KEY }
@@ -88,6 +90,13 @@ class BfxWebsocketClient(object):
             if chanId in bucket.chanIds.keys():
                 await bucket._unsubscribe(chanId=chanId)
 
+    async def close(self, code=1000, reason=str()):
+        if self.websocket != None and self.websocket.open == True:
+            await self.websocket.close(code=code, reason=reason)
+
+        for bucket in self.buckets:
+            await bucket.close(code=code, reason=reason)
+
     def __require_websocket_authentication(function):
         @_require_websocket_connection
         async def wrapper(self, *args, **kwargs):
@@ -131,11 +140,12 @@ class _BfxWebsocketBucket(object):
         self.handler = PublicChannelsHandler(event_emitter=self.event_emitter)
 
     async def _connect(self, index):
-        async with websockets.connect(self.host) as websocket:
+        async for websocket in websockets.connect(self.host):
             self.websocket = websocket
 
             self.__bucket_open_signal(index)
 
+            try:
                 async for message in websocket:
                     message = json.loads(message)
 
@@ -153,6 +163,7 @@ class _BfxWebsocketBucket(object):
                         self.event_emitter.emit("wss-error", message["code"], message["msg"])
                     elif isinstance(message, list) and (chanId := message[0]) and message[1] != HEARTBEAT:
                         self.handler.handle(self.chanIds[chanId], *message[1:])
+            except websockets.ConnectionClosedError: continue
 
     @_require_websocket_connection
     async def _subscribe(self, channel, subId=None, **kwargs):
@@ -177,3 +188,7 @@ class _BfxWebsocketBucket(object):
             "event": "unsubscribe",
             "chanId": chanId
         }))
+
+    @_require_websocket_connection
+    async def close(self, code=1000, reason=str()):
+        await self.websocket.close(code=code, reason=reason)