graphql-transport-ws implementation with `QWebSocket`

pyside_gql_subscriber.py

import uuid
from dataclasses import dataclass, asdict, field, is_dataclass
import json

from qtpy import QtCore as qtc, QtWebSockets as qtws, QtNetwork as qtn
from typing import NamedTuple, Optional, NewType, Union, Any

from tzvui.network.client import HOST, PORT
from tzvui.qthacks import slot

UNSET = NewType('UNSET', None)


class PROTOCOL:
    """
    The WebSocket sub-protocol for this specification is: graphql-transport-ws.
    https://github.com/enisdenjo/graphql-ws/blob/master/PROTOCOL.md
    """
    CONNECTION_INIT = 'connection_init'
    CONNECTION_ACK = 'connection_ack'
    NEXT = 'next'
    ERROR = 'error'
    COMPLETE = 'complete'
    PING = 'ping'
    PONG = 'pong'
    SUBSCRIBE = 'subscribe'

    @classmethod
    def __dict__(cls) -> dict:
        ret = {}
        for k, v in cls.__dict__.items():
            k: str
            if k.isupper():
                ret[k] = v

        return ret


@dataclass(slots=True)
class BaseGqlMessage:
    type: str
    payload: dict = UNSET

    def asdict(self) -> dict:
        data = asdict(self)
        if getattr(self, "payload", None) is UNSET:
            # Unset fields must have a JSON value of "undefined" not "null"
            data.pop("payload")
        return data


@dataclass(slots=True, kw_only=True)
class SubscribeMessage(BaseGqlMessage):
    payload: 'SubscribePayload' = UNSET
    type: str = PROTOCOL.SUBSCRIBE
    id: uuid.UUID = field(default_factory=lambda: uuid.uuid4())


@dataclass(slots=True)
class SubscribePayload:
    query: str
    operationName: str = "subscription"


class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, uuid.UUID):
            return str(obj)
        elif isinstance(obj, BaseGqlMessage):
            return obj.asdict()
        elif is_dataclass(obj):
            return asdict(obj)

class MESSAGES:
    CONNECTION_INIT = BaseGqlMessage(PROTOCOL.CONNECTION_INIT)
    PING = BaseGqlMessage(PROTOCOL.PING)
    PONG = BaseGqlMessage(PROTOCOL.PONG)


class GqlWsTransport(qtws.QWebSocket):
    SUB_PROTOCOL = 'graphql-transport-ws'

    def __init__(self,
                 subscription: SubscribeMessage,
                 parent=None,
                 ping_interval=1000,
                 ping_timeout=500
                 ):
        super().__init__(parent)
        self.subscription = subscription
        self.ping_timer = qtc.QTimer()
        self.ping_timer.setInterval(ping_interval)
        self.ping_timer.timeout.connect(self.on_ping_timeout)
        self.ping_tester_timer = qtc.QTimer()
        self.ping_tester_timer.setInterval(ping_timeout)
        self.ping_tester_timer.timeout.connect(self.on_ping_tester_timeout)
        self.ws_options = qtws.QWebSocketHandshakeOptions()
        self.ws_options.setSubprotocols([self.SUB_PROTOCOL, ])
        self.textMessageReceived.connect(self.on_text_message)
        self.pong.connect(self.on_pong)
        self.connected.connect(self.on_connected)
        self.disconnected.connect(self.on_disconnected)
        self.error.connect(self.on_error)

    @staticmethod
    def dumps(data: Any) -> str:
        return json.dumps(data, cls=MyEncoder)

    @slot
    def open(self, url: str = f"ws://{HOST}:{PORT}/graphql") -> None:
        req = qtn.QNetworkRequest()
        req.setUrl(qtc.QUrl(url))
        super().open(req, self.ws_options)

    def close(self, closeCode=..., reason=...) -> None:
        print("Gracefully closed", closeCode, reason)

    def on_connected(self):
        print("Connected")
        self.sendTextMessage(self.dumps(MESSAGES.CONNECTION_INIT))

    def on_text_message(self, message: str) -> None:
        message = BaseGqlMessage(**json.loads(message))
        match message.type:
            case PROTOCOL.CONNECTION_ACK:
                self.sendTextMessage(self.dumps(MESSAGES.PING))
                self.ping_timer.start()
                self.ping_tester_timer.start()
                self.run_subscription()

            case PROTOCOL.NEXT:
                self.on_gql_next(message)

            case PROTOCOL.PONG:
                self.on_gql_pong()

            case PROTOCOL.PING:
                self.on_gql_ping()

            case PROTOCOL.ERROR:
                self.on_gql_error(message)

            case PROTOCOL.COMPLETE:
                self.on_gql_complete(message)



    def on_error(self, error: qtn.QAbstractSocket.SocketError):
        print(error)

    def run_subscription(self):
        self.sendTextMessage(self.dumps(self.subscription))

    def on_ping_timeout(self):
        self.sendTextMessage(self.dumps(MESSAGES.PING))
        self.ping_tester_timer.start()

    def on_ping_tester_timeout(self):
        raise TimeoutError("Ping timeout reached!")

    def on_gql_pong(self):
        print("Pong received")
        self.ping_timer.stop()
        self.ping_tester_timer.stop()

    def on_gql_ping(self):
        print("Ping received")
        self.sendTextMessage(self.dumps(MESSAGES.PONG))

    def on_gql_error(self, message: BaseGqlMessage):
        print(message)

    def on_gql_complete(self, message: BaseGqlMessage):
        print(message)

    def on_gql_next(self, message: BaseGqlMessage):
        print(message)

if __name__ == "__main__":
    app = qtc.QCoreApplication()
    message = SubscribeMessage(
        payload=SubscribePayload(
            query="subscription subscription {\n  alerts {\n    uuid\n    name\n \n  }\n}")
    )
    client = GqlWsTransport(message)
    client.open()
    app.exec()