"""Connect to pubnub.""" # Taken from https://github.com/bdraco/yalexs/blob/main/yalexs/pubnub_async.py import asyncio import logging import secrets from functools import partial from typing import Callable from pubnub.callbacks import SubscribeCallback from pubnub.enums import PNReconnectionPolicy, PNStatusCategory from pubnub.models.consumer.common import PNStatus from pubnub.models.consumer.pubsub import PNMessageResult from pubnub.pnconfiguration import PNConfiguration from pubnub.pubnub_asyncio import PubNubAsyncio from .containers import SnooData _LOGGER = logging.getLogger(__name__) logging.basicConfig() _LOGGER.setLevel(logging.DEBUG) SHOULD_RECONNECT_CATEGORIES = { PNStatusCategory.PNUnknownCategory, PNStatusCategory.PNUnexpectedDisconnectCategory, PNStatusCategory.PNNetworkIssuesCategory, PNStatusCategory.PNTimeoutCategory, } class SnooPubNub(SubscribeCallback): def __init__(self, pubnub: PubNubAsyncio, device_id: str) -> None: """Initialize the SnooPubNub.""" super().__init__() self.pubnub = pubnub self.device_id = device_id self._subscriptions: set[Callable[[SnooData], None]] = set() # correct type hint self.connected = False self.task: asyncio.Task | None = None def update_token(self, token: str): self.pubnub.config.auth_key = token if self.task: self.task.cancel() # cancel the task if it exists. self.task = None asyncio.create_task(self.run()) # restart the task def presence(self, pubnub: PubNubAsyncio, presence): _LOGGER.debug("Received new presence: %s", presence) def status(self, pubnub: PubNubAsyncio, status: PNStatus) -> None: if not pubnub: self.connected = False return _LOGGER.debug( "Received new status: category=%s error_data=%s error=%s status_code=%s operation=%s", status.category, status.error_data, status.error, status.status_code, status.operation, ) if status.category in SHOULD_RECONNECT_CATEGORIES: self.connected = False if self.task: self.task.cancel() # cancel the task if it exists. self.task = None asyncio.create_task(self.run()) # restart the task elif status.category == PNStatusCategory.PNReconnectedCategory: self.connected = True elif status.category == PNStatusCategory.PNConnectedCategory: self.connected = True def message(self, pubnub: PubNubAsyncio, message: PNMessageResult) -> None: # Handle new messages _LOGGER.debug( "Received new messages on channel %s for device_id: %s with timetoken: %s: %s", message.channel, self.device_id, message.timetoken, message.message, ) if message.channel.split(".")[0] == "ActivityState": if "system_state" in message.message: for callback in self._subscriptions: data = SnooData.from_dict(message.message) _LOGGER.debug(data) callback(data) def subscribe(self, update_callback: Callable[[SnooData], None]) -> Callable[[], None]: """Add an callback subscriber. Returns a callable that can be used to unsubscribe. """ self._subscriptions.add(update_callback) return partial(self._unsubscribe, update_callback) def _unsubscribe(self, update_callback: Callable[[SnooData], None]) -> None: self._subscriptions.remove(update_callback) async def run(self) -> None: """Run the pubnub loop.""" if self.task: # prevent multiple tasks from running return pnconfig = PNConfiguration() pnconfig.subscribe_key = "sub-c-97bade2a-483d-11e6-8b3b-02ee2ddab7fe" pnconfig.publish_key = "pub-c-699074b0-7664-4be2-abf8-dcbb9b6cd2bf" pnconfig.user_id = secrets.token_urlsafe(16) pnconfig.auth_key = self.pubnub.config.auth_key pnconfig.reconnect_policy = PNReconnectionPolicy.EXPONENTIAL self.pubnub.pnconfig = pnconfig self.pubnub.add_listener(self) self.pubnub.subscribe().channels( [f"ActivityState.{self.device_id}", f"ControlCommand.{self.device_id}"] ).execute()