fix message typing in channel manager, minor improvements

This commit is contained in:
Timothy Pogue 2022-10-23 11:42:59 -06:00
parent 9cffa3ca2b
commit 94b65a007a
3 changed files with 17 additions and 19 deletions

View File

@ -73,7 +73,7 @@ async def _process_consumer_request(
whitelist = rpc._ws_request_whitelist() whitelist = rpc._ws_request_whitelist()
# Format response # Format response
response = WSWhitelistMessage(data=whitelist).dict(exclude_none=True) response = WSWhitelistMessage(data=whitelist)
# Send it back # Send it back
await channel_manager.send_direct(channel, response) await channel_manager.send_direct(channel, response)
@ -89,7 +89,7 @@ async def _process_consumer_request(
# For every dataframe, send as a separate message # For every dataframe, send as a separate message
for _, message in analyzed_df.items(): for _, message in analyzed_df.items():
response = WSAnalyzedDFMessage(data=message).dict(exclude_none=True) response = WSAnalyzedDFMessage(data=message)
await channel_manager.send_direct(channel, response) await channel_manager.send_direct(channel, response)

View File

@ -10,6 +10,7 @@ from freqtrade.rpc.api_server.ws.proxy import WebSocketProxy
from freqtrade.rpc.api_server.ws.serializer import (HybridJSONWebSocketSerializer, from freqtrade.rpc.api_server.ws.serializer import (HybridJSONWebSocketSerializer,
WebSocketSerializer) WebSocketSerializer)
from freqtrade.rpc.api_server.ws.types import WebSocketType from freqtrade.rpc.api_server.ws.types import WebSocketType
from freqtrade.rpc.api_server.ws_schemas import WSMessageSchema
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@ -54,8 +55,8 @@ class WebSocketChannel:
return f"WebSocketChannel({self.channel_id}, {self.remote_addr})" return f"WebSocketChannel({self.channel_id}, {self.remote_addr})"
@property @property
def raw(self): def raw_websocket(self):
return self._websocket.raw return self._websocket.raw_websocket
@property @property
def remote_addr(self): def remote_addr(self):
@ -192,29 +193,26 @@ class ChannelManager:
for websocket in self.channels.copy().keys(): for websocket in self.channels.copy().keys():
await self.on_disconnect(websocket) await self.on_disconnect(websocket)
self.channels = dict() async def broadcast(self, message: WSMessageSchema):
async def broadcast(self, data):
""" """
Broadcast data on all Channels Broadcast a message on all Channels
:param data: The data to send :param message: The message to send
""" """
with self._lock: with self._lock:
message_type = data.get('type')
for channel in self.channels.copy().values(): for channel in self.channels.copy().values():
if channel.subscribed_to(message_type): if channel.subscribed_to(message.type):
await self.send_direct(channel, data) await self.send_direct(channel, message)
async def send_direct(self, channel, data): async def send_direct(self, channel: WebSocketChannel, message: WSMessageSchema):
""" """
Send data directly through direct_channel only Send a message directly through direct_channel only
:param direct_channel: The WebSocketChannel object to send data through :param direct_channel: The WebSocketChannel object to send the message through
:param data: The data to send :param message: The message to send
""" """
if not await channel.send(data): if not await channel.send(message.dict(exclude_none=True)):
await self.on_disconnect(channel.raw) await self.on_disconnect(channel.raw_websocket)
def has_channels(self): def has_channels(self):
""" """

View File

@ -16,7 +16,7 @@ class WebSocketProxy:
self._websocket: Union[FastAPIWebSocket, WebSocket] = websocket self._websocket: Union[FastAPIWebSocket, WebSocket] = websocket
@property @property
def raw(self): def raw_websocket(self):
return self._websocket return self._websocket
@property @property