From 5dc69cc655a621a42fa42f7b10384d11a1794f9d Mon Sep 17 00:00:00 2001
From: onyx-and-iris <code@onyxandiris.online>
Date: Wed, 28 Jun 2023 02:45:04 +0100
Subject: [PATCH] lazy load tomllib/tomli

raise SteamlabsSIOConnectionError on connection error

Path.home() / ".config" / "streamlabsio" / "config.toml"
added to config filepaths

raw kwarg added.

module level logger added.
---
 streamlabsio/client.py | 58 +++++++++++++++++++++++++++++++-----------
 streamlabsio/error.py  |  2 ++
 2 files changed, 45 insertions(+), 15 deletions(-)
 create mode 100644 streamlabsio/error.py

diff --git a/streamlabsio/client.py b/streamlabsio/client.py
index e554066..189b53a 100644
--- a/streamlabsio/client.py
+++ b/streamlabsio/client.py
@@ -1,22 +1,21 @@
 import logging
 from pathlib import Path
+from typing import Optional
 
 import socketio
 from observable import Observable
 
-try:
-    import tomllib
-except ModuleNotFoundError:
-    import tomli as tomllib
-
+from .error import SteamlabsSIOConnectionError
 from .models import as_dataclass
 
+logger = logging.getLogger(__name__)
+
 
 class Client:
-    logger = logging.getLogger("socketio.client")
-
-    def __init__(self, token=None):
+    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)
@@ -27,15 +26,43 @@ class Client:
         self.youtube = ("follow", "subscription", "superchat")
 
     def __enter__(self):
-        self.sio.connect(f"https://sockets.streamlabs.com?token={self.token}")
+        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:
-        filepath = Path.cwd() / "config.toml"
-        with open(filepath, "rb") as f:
-            conn = tomllib.load(f)
-            assert "token" in conn.get("streamlabs")
-        return conn["streamlabs"].get("token")
+        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")
@@ -47,8 +74,9 @@ class Client:
             self.obs.trigger(
                 data["for"],
                 data["type"],
-                as_dataclass(data["type"], *data["message"]),
+                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")
diff --git a/streamlabsio/error.py b/streamlabsio/error.py
new file mode 100644
index 0000000..8cc512b
--- /dev/null
+++ b/streamlabsio/error.py
@@ -0,0 +1,2 @@
+class SteamlabsSIOConnectionError(Exception):
+    """Exception raised when connection errors occur"""