pigway-frame-flow/pigway_frame_flow/core/registry.py

import yaml

from pigway_frame_flow.drivers.base import FrameDeviceDriver
from pigway_frame_flow.drivers.kokonna.client import KokonnaClient
from pigway_frame_flow.drivers.kokonna.driver import KokonnaDriver
from pigway_frame_flow.drivers.kokonna.mock import MockClient
from pigway_frame_flow.drivers.kokonna.widgets import WidgetReader
from pigway_frame_flow.errors import ErrorCode as E
from pigway_frame_flow.errors import FlowError
from pigway_frame_flow.settings import secret


class DeviceRegistry:
    def __init__(self, settings):
        self.settings = settings
        config = settings.config_dir / "devices.yml"
        devices = (
            yaml.safe_load(config.read_text())["devices"]
            if config.exists()
            else {
                settings.device_id: {
                    "driver": "kokonna",
                    "api_key_env": "KOKONNA_API_KEY",
                    "enabled": True,
                }
            }
        )
        self.drivers: dict[str, FrameDeviceDriver] = {}
        self.keys = []
        for id, entry in devices.items():
            if not entry.get("enabled", True):
                continue
            if entry.get("driver") != "kokonna":
                raise ValueError("Unsupported configured device driver")
            if not __import__("re").fullmatch(r"[a-z0-9-]{1,60}", id):
                raise ValueError("Invalid configured device ID")
            client: MockClient | KokonnaClient
            if settings.mock:
                root = settings.data_dir / "cache" / id
                root.mkdir(parents=True, exist_ok=True)
                client = MockClient(root)
            else:
                key = (
                    settings.api_key
                    if entry.get("api_key_env", "KOKONNA_API_KEY") == "KOKONNA_API_KEY"
                    else secret(str(entry["api_key_env"]))
                )
                self.keys.append(key)
                settings.additional_secrets.append(key)
                client = KokonnaClient(key, settings.write_enabled)
            widget_token = secret(
                str(entry.get("widget_token_env", "KOKONNA_WIDGET_TOKEN"))
            )
            settings.additional_secrets.append(widget_token)
            widgets = WidgetReader(
                widget_token, entry.get("widget_robot_id"), settings.mock
            )
            self.drivers[id] = KokonnaDriver(client, id, settings.sanitize, widgets)

    def get(self, id: str) -> FrameDeviceDriver:
        if id not in self.drivers:
            raise FlowError(E.DEVICE_NOT_FOUND, "Device not found", 404)
        return self.drivers[id]