freqtrade_origin/freqtrade/rpc/replicate/proxy.py

from typing import TYPE_CHECKING, Union

from fastapi import WebSocket as FastAPIWebSocket
from websockets import WebSocketClientProtocol as WebSocket


if TYPE_CHECKING:
    from freqtrade.rpc.replicate.types import WebSocketType


class WebSocketProxy:
    """
    WebSocketProxy object to bring the FastAPIWebSocket and websockets.WebSocketClientProtocol
    under the same API
    """

    def __init__(self, websocket: WebSocketType):
        self._websocket: Union[FastAPIWebSocket, WebSocket] = websocket

    async def send(self, data):
        """
        Send data on the wrapped websocket
        """
        if hasattr(self._websocket, "send_bytes"):
            await self._websocket.send_bytes(data)
        else:
            await self._websocket.send(data)

    async def recv(self):
        """
        Receive data on the wrapped websocket
        """
        if hasattr(self._websocket, "receive_bytes"):
            return await self._websocket.receive_bytes()
        else:
            return await self._websocket.recv()

    async def ping(self):
        """
        Ping the websocket, not supported by FastAPI WebSockets
        """
        if hasattr(self._websocket, "ping"):
            return await self._websocket.ping()
        return False

    async def close(self, code: int = 1000):
        """
        Close the websocket connection, only supported by FastAPI WebSockets
        """
        if hasattr(self._websocket, "close"):
            return await self._websocket.close(code)
        pass

    async def accept(self):
        """
        Accept the WebSocket connection, only support by FastAPI WebSockets
        """
        if hasattr(self._websocket, "accept"):
            return await self._websocket.accept()
        pass