zapsaber/datapuller.py

import asyncio
import enum
import json
import logging

import websockets


logger = logging.getLogger(__name__)


class BaseData:
    def __init__(self, host: str = '127.0.0.1', port: int = 2946,
                 path: str = '/BSDataPuller'):
        self.host = host
        self.port = port
        self.path = path

    async def listen(self):
        try:
            async with websockets.connect(
                    'ws://{}:{}{}'.format(self.host, self.port, self.path)) as websocket:
                async for message in websocket:
                    logger.debug(message)
                    if message:
                        self._handle(json.loads(message))
        except Exception as e:
            logger.critical(e)
            raise e


class LiveData(BaseData):
    TRIGGERS = enum.Enum('Triggers',
            ['Unknown', 'TimerElapsed', 'NoteMissed', 'EnergyChange', 'ScoreChange'])
    TRIGGERS_INVERSE = [
        TRIGGERS.Unknown,
        TRIGGERS.TimerElapsed,
        TRIGGERS.NoteMissed,
        TRIGGERS.EnergyChange,
        TRIGGERS.ScoreChange,
    ]

    def __init__(self, host: str = '127.0.0.1', port: int = 2946,
                 path: str = '/BSDataPuller/LiveData'):
        super().__init__(host, port, path)
        self.events = {}

    def on(self, event: str, function):
        if event not in self.events:
            self.events[event] = []
        self.events[event].append(function)

    def _handle(self, message: str):
        invert = self.TRIGGERS_INVERSE[message['EventTrigger']].name
        if invert in self.events:
            for event in self.events[invert]:
                event(message)

        
class MapData(BaseData):
    def __init__(self, host: str = '127.0.0.1', port: int = 2946,
                 path: str = '/BSDataPuller/MapData'):
        super().__init__(host, port, path)

    def on_update(self, function):
        self.function = function

    def _handle(self, message: str):
        self.function(message)