diff --git a/constants.py b/constants.py index e193476..df1d0ac 100644 --- a/constants.py +++ b/constants.py @@ -2,11 +2,11 @@ from __future__ import annotations from copy import copy from datetime import timedelta -from typing import Any, Optional, Union, Dict, Callable +from typing import Any, Optional, Dict, Literal, Callable # Typing JsonType = Dict[str, Any] -TopicProcess = Callable[[JsonType], Any] +TopicProcess = Callable[[int, JsonType], Any] # Values MAX_WEBSOCKETS = 8 WS_TOPICS_LIMIT = 50 @@ -130,25 +130,34 @@ GQL_OPERATIONS: Dict[str, GQLOperation] = { class WebsocketTopic: def __init__( self, - topic_category: str, + category: Literal["User", "Channel"], topic_name: str, - target_id: Union[int, str], + target_id: int, process: TopicProcess, ): - self._id: str = f"{WEBSOCKET_TOPICS[topic_category][topic_name]}.{target_id}" + self._id: str = f"{WEBSOCKET_TOPICS[category][topic_name]}.{target_id}" + assert isinstance(target_id, int) + self._target_id = target_id self._process: TopicProcess = process - def __call__(self, *args): - return self._process(*args) + @classmethod + def as_str(cls, category: Literal["User", "Channel"], topic_name: str, target_id: int) -> str: + return f"{WEBSOCKET_TOPICS[category][topic_name]}.{target_id}" + + def __call__(self, message: JsonType): + return self._process(self._target_id, message) def __str__(self) -> str: return self._id + def __repr__(self) -> str: + return f"Topic({self._id})" + def __eq__(self, other): - if isinstance(other, str): - return self._id == other - elif isinstance(other, WebsocketTopic): + if isinstance(other, WebsocketTopic): return self._id == other._id + elif isinstance(other, str): + return self._id == other return NotImplemented def __hash__(self) -> int: diff --git a/twitch.py b/twitch.py index 0e4547f..30733c4 100644 --- a/twitch.py +++ b/twitch.py @@ -133,12 +133,7 @@ class Twitch: self.channels[channel.id] = channel # Sub to these channel updates topics: List[WebsocketTopic] = [ - WebsocketTopic( - "Channel", - "VideoPlayback", - channel_id, - partial(self.process_stream_state, channel_id), - ) + WebsocketTopic("Channel", "VideoPlayback", channel_id, self.process_stream_state) for channel_id in self.channels ] self.websocket.add_topics(topics) diff --git a/websocket.py b/websocket.py index 5aaac5e..6cb6b0e 100644 --- a/websocket.py +++ b/websocket.py @@ -247,14 +247,14 @@ class Websocket: self.topics[str(topic)] = topic self._topics_changed.set() - def remove_topics(self, topics_set: Set[WebsocketTopic]): - existing = topics_set.intersection(self.topics.values()) + def remove_topics(self, topics_set: Set[str]): + existing = topics_set.intersection(self.topics.keys()) if not existing: # nothing to remove from here return topics_set.difference_update(existing) for topic in existing: - del self.topics[str(topic)] + del self.topics[topic] self._topics_changed.set() async def send(self, message: JsonType): @@ -281,6 +281,8 @@ class WebsocketPool: async def start(self): await self._twitch.wait_until_login() + if self.running: + return # Add default topics assert self._twitch._user_id is not None user_id = self._twitch._user_id @@ -301,7 +303,7 @@ class WebsocketPool: if not topics_set: # nothing to add return - topics_set.difference_update(ws.topics.values() for ws in self.websockets) + topics_set.difference_update(*(ws.topics.values() for ws in self.websockets)) if not topics_set: # none left to add return @@ -323,7 +325,7 @@ class WebsocketPool: # if we're here, there were leftover topics after filling up all websockets raise MinerException("Maximum topics limit has been reached") - def remove_topics(self, topics: Iterable[WebsocketTopic]): + def remove_topics(self, topics: Iterable[str]): topics_set = set(topics) if not topics_set: # nothing to remove @@ -332,20 +334,20 @@ class WebsocketPool: ws.remove_topics(topics_set) # count up all the topics - if we happen to have more websockets connected than needed, # stop the last one and recycle topics from it - repeat until we have enough - topics = [] + recycled_topics: List[WebsocketTopic] = [] while True: count = sum(len(ws.topics) for ws in self.websockets) if count <= (len(self.websockets) - 1) * WS_TOPICS_LIMIT: ws = self.websockets.pop() - topics.extend(ws.topics.values()) + recycled_topics.extend(ws.topics.values()) ws.stop_nowait() else: break - if topics: - self.add_topics(topics) + if recycled_topics: + self.add_topics(recycled_topics) @task_wrapper - async def process_drops(self, message: JsonType): + async def process_drops(self, user_id: int, message: JsonType): drop_id = message["data"]["drop_id"] drop: Optional[TimedDrop] = None for campaign in self._twitch.inventory: @@ -373,7 +375,7 @@ class WebsocketPool: self._twitch.reevaluate_campaigns() @task_wrapper - async def process_points(self, message: JsonType): + async def process_points(self, user_id: int, message: JsonType): msg_type = message["type"] if msg_type == "points-earned": points = message["data"]["point_gain"]["total_points"]