streamlabs-socketio-py/streamlabsio/client.py

import logging
from pathlib import Path
from typing import Optional

import socketio
from observable import Observable

from .error import SteamlabsSIOConnectionError
from .models import as_dataclass

logger = logging.getLogger(__name__)


class Client:
    def __init__(self, token=None, raw=False):
        self.logger = logger.getChild(self.__class__.__name__)
        self.token = token or self._token_from_toml()
        self._raw = raw
        self.sio = socketio.Client()
        self.sio.on("connect", self.connect_handler)
        self.sio.on("event", self.event_handler)
        self.sio.on("disconnect", self.disconnect_handler)
        self.obs = Observable()
        self.streamlabs = ("donation",)
        self.twitch = ("follow", "subscription", "host", "bits", "raids")
        self.youtube = ("follow", "subscription", "superchat")

    def __enter__(self):
        try:
            self.sio.connect(f"https://sockets.streamlabs.com?token={self.token}")
        except socketio.exceptions.ConnectionError as e:
            self.logger.exception(f"{type(e).__name__}: {e}")
            raise SteamlabsSIOConnectionError(
                "no connection could be established to the Streamlabs SIO server"
            ) from e
        return self

    def _token_from_toml(self) -> str:
        try:
            import tomllib
        except ModuleNotFoundError:
            import tomli as tomllib

        def get_filepath() -> Optional[Path]:
            filepaths = (
                Path.cwd() / "config.toml",
                Path.home() / ".config" / "streamlabsio" / "config.toml",
            )
            for filepath in filepaths:
                if filepath.exists():
                    return filepath

        try:
            filepath = get_filepath()
            if not filepath:
                raise FileNotFoundError("config.toml was not found")
            with open(filepath, "rb") as f:
                conn = tomllib.load(f)
                assert "token" in conn.get(
                    "streamlabs"
                ), "token not found in config.toml"
            return conn["streamlabs"].get("token")
        except (FileNotFoundError, tomllib.TOMLDecodeError) as e:
            self.logger.error(f"{type(e).__name__}: {e}")
            raise

    def connect_handler(self):
        self.logger.info("Connected to Streamlabs Socket API")

    def event_handler(self, data):
        if "for" in data and data["type"] in set(
            self.streamlabs + self.twitch + self.youtube
        ):
            self.obs.trigger(
                data["for"],
                data["type"],
                data if self._raw else as_dataclass(data["type"], *data["message"]),
            )
            self.logger.debug(data)

    def disconnect_handler(self):
        self.logger.info("Disconnected from Streamlabs Socket API")

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.sio.disconnect()


def connect(**kwargs):
    SIO_cls = Client
    return SIO_cls(**kwargs)