diff --git a/litecord/gateway/utils.py b/litecord/gateway/utils.py new file mode 100644 index 0000000..d52d7df --- /dev/null +++ b/litecord/gateway/utils.py @@ -0,0 +1,30 @@ +""" + +Litecord +Copyright (C) 2018-2019 Luna Mendes + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU General Public License as published by +the Free Software Foundation, version 3 of the License. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU General Public License for more details. + +You should have received a copy of the GNU General Public License +along with this program. If not, see . + +""" + +import asyncio + +class WebsocketFileHandler: + """A handler around a websocket that wraps normal I/O calls into + the websocket's respective asyncio calls via asyncio.ensure_future.""" + def __init__(self, ws): + self.ws = ws + + def write(self, data): + """Write data into the websocket""" + asyncio.ensure_future(self.ws.send(data)) diff --git a/litecord/gateway/websocket.py b/litecord/gateway/websocket.py index eea81a8..365163c 100644 --- a/litecord/gateway/websocket.py +++ b/litecord/gateway/websocket.py @@ -49,6 +49,8 @@ from litecord.gateway.encoding import ( encode_json, decode_json, encode_etf, decode_etf ) +from litecord.gateway.utils import WebsocketFileHandler + from litecord.storage import int_ log = Logger(__name__) @@ -155,7 +157,11 @@ class GatewayWebsocket: await self._chunked_send(data2, 1024) async def _zstd_stream_send(self, encoded): - pass + compressor = self.wsp.zsctx.stream_writer( + WebsocketFileHandler(self.ws)) + + compressor.write(encoded) + compressor.flush(zstd.FLUSH_FRAME) async def send(self, payload: Dict[str, Any]): """Send a payload to the websocket.