diff --git a/README.md b/README.md index ccba3c5..5a9a69f 100644 --- a/README.md +++ b/README.md @@ -156,12 +156,12 @@ Both belong to the TikTokLive `Event` type and can be listened to. The following - `ConnectEvent` - Triggered when the Webcast connection is initiated - `DisconnectEvent` - Triggered when the Webcast connection closes (including the livestream ending) - `LiveEndEvent` - Triggered when the livestream ends -- `LivePausedEvent` - Triggered when the livestream is paused -- `LiveUnpausedEvent` - Triggered when the livestream is unpaused +- `LivePauseEvent` - Triggered when the livestream is paused +- `LiveUnpauseEvent` - Triggered when the livestream is unpaused - `FollowEvent` - Triggered when a user in the livestream follows the streamer - `ShareEvent` - Triggered when a user shares the livestream -- `UnknownEvent` - Events not currently tracked by TikTokLive as they have not been reverse-engineered - +- `WebsocketResponseEvent` - Triggered when any event is received (contains the event) +- `UnknownEvent` - An instance of `WebsocketResponseEvent` thrown whenever an event does not have an existing definition, useful for debugging ### Proto Events diff --git a/TikTokLive/client/client.py b/TikTokLive/client/client.py index 227fe49..0ca9c29 100644 --- a/TikTokLive/client/client.py +++ b/TikTokLive/client/client.py @@ -16,13 +16,17 @@ from TikTokLive.client.web.web_settings import WebDefaults from TikTokLive.client.ws.ws_client import WebcastWSClient from TikTokLive.events import Event, EventHandler -from TikTokLive.events.custom_events import UnknownEvent, ConnectEvent, FollowEvent, ShareEvent, LiveEndEvent, \ - DisconnectEvent, LivePausedEvent, LiveUnpausedEvent +from TikTokLive.events.custom_events import WebsocketResponseEvent, ConnectEvent, FollowEvent, ShareEvent, LiveEndEvent, \ + DisconnectEvent, LivePauseEvent, LiveUnpauseEvent, UnknownEvent, CustomEvent from TikTokLive.events.proto_events import EVENT_MAPPINGS, ProtoEvent, ControlEvent from TikTokLive.proto import WebcastResponse, WebcastResponseMessage, ControlAction class TikTokLiveClient(AsyncIOEventEmitter): + """ + A client to connect to & read from TikTok LIVE streams + + """ def __init__( self, @@ -32,6 +36,17 @@ def __init__( web_kwargs: dict = {}, ws_kwargs: dict = {} ): + """ + Instantiate the TikTokLiveClient client + + :param unique_id: The username of the creator to connect to + :param web_proxy: An optional proxy used for HTTP requests + :param ws_proxy: An optional proxy used for the WebSocket connection + :param web_kwargs: Optional arguments used by the HTTP client + :param ws_kwargs: Optional arguments used by the WebSocket client + + """ + super().__init__() unique_id = self.parse_unique_id(unique_id) @@ -176,7 +191,13 @@ async def disconnect(self) -> None: self._event_loop_task = None async def _client_loop(self, initial_response: WebcastResponse) -> None: - """Run the main client loop to handle events""" + """ + Run the main client loop to handle events + + :param initial_response: The WebcastResponse retrieved from the sign server with connection info + :return: None + + """ async for event in self._ws_loop(initial_response): @@ -191,7 +212,13 @@ async def _client_loop(self, initial_response: WebcastResponse) -> None: self.emit(ev.type, ev) async def _ws_loop(self, initial_response: WebcastResponse) -> AsyncIterator[Optional[Event]]: - """Run the websocket loop to handle incoming WS messages""" + """ + Run the websocket loop to handle incoming WS messages + + :param initial_response: The WebcastResponse retrieved from the sign server with connection info + :return: None + + """ first_event: bool = True @@ -213,7 +240,13 @@ async def _ws_loop(self, initial_response: WebcastResponse) -> AsyncIterator[Opt yield event def _build_connect_info(self, initial_response: WebcastResponse) -> Tuple[str, dict]: - """Create connection info for starting the connection""" + """ + Create connection info for starting the connection + + :param initial_response: The WebcastResponse retrieved from the sign server with connection info + :return: None + + """ connect_uri: str = ( initial_response.push_server @@ -228,60 +261,133 @@ def _build_connect_info(self, initial_response: WebcastResponse) -> Tuple[str, d return connect_uri, connect_headers def on(self, event: Type[Event], f: Optional[EventHandler] = None) -> Union[Handler, Callable[[Handler], Handler]]: + """ + Decorator that can be used to register a Python function as an event listener + + :param event: The event to listen to + :param f: The function to handle the event + :return: The wrapped function as a generated `pyee.Handler` object + + """ + return super(TikTokLiveClient, self).on(event.get_type(), f) def add_listener(self, event: Type[Event], f: EventHandler) -> Handler: + """ + Method that can be used to register a Python function as an event listener + + :param event: The event to listen to + :param f: The function to handle the event + :return: The generated `pyee.Handler` object + + """ + return super().add_listener(event=event.get_type(), f=f) - def has_listener(self, *events: Type[Event]) -> bool: - return any(event.__name__ in self._events for event in events) + def has_listener(self, event: Type[Event]) -> bool: + """ + Check whether the client is listening to a given event + + :param event: The event to check listening for + :return: Whether it is being listened to + + """ + + return event.__name__ in self._events def _parse_webcast_response(self, response: Optional[WebcastResponseMessage]) -> List[Event]: - """Parse incoming webcast responses""" + """ + Parse incoming webcast responses into events that can be emitted + :param response: The WebcastResponseMessage protobuf message + :return: A list of events that can be gleamed from this event + + """ + + # Invalid response handler if response is None: - self._logger.warning("Received a null response from the Webcast server.") + self._logger.warning("Received a null WebcastResponseMessage from the Webcast server.") return [] + # Get the proto mapping for proto-events event_type: Optional[Type[ProtoEvent]] = EVENT_MAPPINGS.get(response.method) + response_event: Event = WebsocketResponseEvent().from_pydict(response.to_dict()) - if not event_type: - return [UnknownEvent().from_pydict(response.to_dict())] + # If the event is not tracked, return + if event_type is None: + return [response_event, UnknownEvent().from_pydict(response.to_dict())] - event: Event = event_type().parse(response.payload) + # Get the underlying events + proto_event: ProtoEvent = event_type().parse(response.payload) + parsed_events: List[Event] = [response_event, proto_event] + custom_event: Optional[Event] = self._parse_custom_event(response, proto_event) - # Handle stream control events - if isinstance(event, ControlEvent): - return_events: List[Event] = [event] + # Add the custom event IF not null + return [custom_event, *parsed_events] if custom_event else parsed_events + @classmethod + def _parse_custom_event(cls, response: WebcastResponseMessage, event: ProtoEvent) -> Optional[CustomEvent]: + """ + Extract CustomEvent events from existing ProtoEvent events + + :param response: The WebcastResponseMessage to parse for the custom event + :param event: The ProtoEvent to parse for the custom event + :return: The event, if one exists + + """ + + # LiveEndEvent, LivePauseEvent, LiveUnpauseEvent + if isinstance(event, ControlEvent): if event.action == ControlAction.STREAM_ENDED: - return_events.append(LiveEndEvent().parse(response.payload)) + return LiveEndEvent().parse(response.payload) elif event.action == ControlAction.STREAM_PAUSED: - return_events.append(LivePausedEvent().parse(response.payload)) + return LivePauseEvent().parse(response.payload) elif event.action == ControlAction.STREAM_PAUSED: - return_events.append(LiveUnpausedEvent().parse(response.payload)) + return LiveUnpauseEvent().parse(response.payload) + return None - return return_events + # FollowEvent + if "follow" in event.common.display_text.key: + return FollowEvent().parse(response.payload) - # Handle follow & share events - if self.has_listener(FollowEvent, ShareEvent): - if "follow" in event.common.display_text.key: - return [FollowEvent().parse(response.payload), event] - if "share" in event.common.display_text.key: - return [ShareEvent().parse(response.payload), event] + # ShareEvent + if "share" in event.common.display_text.key: + return ShareEvent().parse(response.payload) - return [event] + # Not a custom event + return None @property def room_id(self) -> Optional[str]: + """ + The room ID the user is currently connected to + + :return: Room ID or None + + """ + return self._room_id @property def web(self) -> TikTokWebClient: + """ + The HTTP client that this client uses for requests + + :return: A copy of the TikTokWebClient + + """ + return self._web @property def _asyncio_loop(self) -> AbstractEventLoop: + """ + Property to return the existing or generate a new asyncio event loop + + :return: An asyncio event loop + + """ + try: return asyncio.get_running_loop() except RuntimeError: @@ -289,16 +395,44 @@ def _asyncio_loop(self) -> AbstractEventLoop: @property def connected(self) -> bool: + """ + Whether the WebSocket client is currently connected to TikTok + + :return: Connection status + + """ + return self._ws.connected @property def logger(self) -> logging.Logger: + """ + The internal logger used by TikTokLive + + :return: An instance of a `logging.Logger` + + """ + return self._logger @property def gift_info(self) -> Optional[dict]: + """ + Information about the stream's gifts *if* fetch_gift_info=True when starting the client e.g. with `client.run`) + + :return: The stream gift info + + """ + return self._gift_info @property def room_info(self) -> Optional[dict]: + """ + Information about the room *if* fetch_room_info=True when starting the client (e.g. with `client.run`) + + :return: Dictionary of room info + + """ + return self._room_info diff --git a/TikTokLive/events/custom_events.py b/TikTokLive/events/custom_events.py index 7ba12dc..0c477ad 100644 --- a/TikTokLive/events/custom_events.py +++ b/TikTokLive/events/custom_events.py @@ -8,9 +8,16 @@ from TikTokLive.proto import WebcastResponseMessage -class UnknownEvent(WebcastResponseMessage, BaseEvent): +class WebsocketResponseEvent(WebcastResponseMessage, BaseEvent): """ - Thrown when a Webcast message is received that is NOT tracked by TikTokLive yet. + Triggered when any event is received from the WebSocket + + """ + + +class UnknownEvent(WebsocketResponseEvent): + """ + Triggered when a Webcast message is received that is NOT tracked by TikTokLive yet. """ @@ -40,14 +47,14 @@ class LiveEndEvent(ControlEvent): """ -class LivePausedEvent(ControlEvent): +class LivePauseEvent(ControlEvent): """ Thrown when the stream is paused """ -class LiveUnpausedEvent(ControlEvent): +class LiveUnpauseEvent(ControlEvent): """ Thrown when a paused stream is unpaused @@ -84,22 +91,26 @@ def users_joined(self) -> Optional[int]: CustomEvent: Type = Union[ + WebsocketResponseEvent, UnknownEvent, ConnectEvent, FollowEvent, ShareEvent, LiveEndEvent, + LivePauseEvent, + LiveUnpauseEvent, DisconnectEvent ] __all__ = [ + "WebsocketResponseEvent", "UnknownEvent", "ConnectEvent", "FollowEvent", "ShareEvent", "LiveEndEvent", - "LivePausedEvent", - "LiveUnpausedEvent", + "LivePauseEvent", + "LiveUnpauseEvent", "CustomEvent", "DisconnectEvent" ]