Improve logging and add hack to view

one short

src/iSponsorBlockTV/api_helpers.py

@@ -99,7 +99,6 @@ async def search_channels(self, channel):
     @AsyncConditionalTTL(time_to_live=300, maxsize=10)  # 5 minutes for non-locked segments
     async def get_segments(self, vid_id):
         if await self.is_whitelisted(vid_id):
-            print("Video is whitelisted")
             return ([], True)  # Return empty list and True to indicate that the cache should last forever
         vid_id_hashed = sha256(vid_id.encode("utf-8")).hexdigest()[
                         :4

src/iSponsorBlockTV/logging_helpers.py

@@ -0,0 +1,113 @@
+import logging
+
+from rich.logging import RichHandler
+from rich._log_render import LogRender
+from rich.text import Text
+from rich.style import Style
+class LogHandler(RichHandler):
+    def __init__(self, device_name, log_name_len, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.filter_strings = []
+        self._log_render = LogRender(
+            device_name = device_name,
+            log_name_len = log_name_len,
+            show_time=True,
+            show_level=True,
+            show_path=False,
+            time_format="[%x %X]",
+            omit_repeated_times=True,
+            level_width=None)
+
+    def add_filter_string(self, s):
+        self.filter_strings.append(s)
+
+    def _filter(self, s):
+        for i in self.filter_strings:
+            s = s.replace(i, "REDACTED")
+        return s
+
+    def format(self, record):
+        original = super().format(record)
+        return self._filter(original)
+
+
+class LogRender(LogRender):
+    def __init__(self, device_name, log_name_len, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.filter_strings = []
+        self.log_name_len = log_name_len
+        self.device_name = device_name
+
+    def __call__(
+        self,
+        console,
+        renderables,
+        log_time,
+        time_format = None,
+        level = "",
+        path = None,
+        line_no = None,
+        link_path = None,
+    ):
+        from rich.containers import Renderables
+        from rich.table import Table
+
+        output = Table.grid(padding=(0, 1))
+        output.expand = True
+        if self.show_time:
+            output.add_column(style="log.time")
+        if self.show_level:
+            output.add_column(style="log.level", width=self.level_width)
+        output.add_column(width = self.log_name_len, style=Style(color="yellow"), overflow="fold")
+        output.add_column(ratio=1, style="log.message", overflow="fold")
+        if self.show_path and path:
+            output.add_column(style="log.path")
+        row = []
+        if self.show_time:
+            log_time = log_time or console.get_datetime()
+            time_format = time_format or self.time_format
+            if callable(time_format):
+                log_time_display = time_format(log_time)
+            else:
+                log_time_display = Text(log_time.strftime(time_format))
+            if log_time_display == self._last_time and self.omit_repeated_times:
+                row.append(Text(" " * len(log_time_display)))
+            else:
+                row.append(log_time_display)
+                self._last_time = log_time_display
+        if self.show_level:
+            row.append(level)
+        row.append(Text(self.device_name))
+        row.append(Renderables(renderables))
+        if self.show_path and path:
+            path_text = Text()
+            path_text.append(
+                path, style=f"link file://{link_path}" if link_path else ""
+            )
+            if line_no:
+                path_text.append(":")
+                path_text.append(
+                    f"{line_no}",
+                    style=f"link file://{link_path}#{line_no}" if link_path else "",
+                )
+            row.append(path_text)
+
+        output.add_row(*row)
+        return output
+
+
+class LogFormatter(logging.Formatter):
+    def __init__(self, fmt=None, datefmt=None, style="%", validate=True, filter_strings=None):
+        super().__init__(fmt, datefmt, style, validate)
+        self.filter_strings = filter_strings or []
+    def add_filter_string(self, s):
+        self.filter_strings.append(s)
+
+    def _filter(self, s):
+        print(s)
+        for i in self.filter_strings:
+            s = s.replace(i, "REDACTED")
+        return s
+    def format(self, record):
+        original = logging.Formatter.format(self, record)
+        return self._filter(original)

src/iSponsorBlockTV/main.py

@@ -2,18 +2,33 @@
 import aiohttp
 import time
 import logging
-from . import api_helpers, ytlounge
+import logging
+import rich
+from rich.logging import RichHandler
+from signal import signal, SIGINT, SIGTERM
+
+from . import api_helpers, ytlounge, logging_helpers
 
 
 class DeviceListener:
-    def __init__(self, api_helper, config, device):
+    def __init__(self, api_helper, config, device, log_name_len, debug: bool):
         self.task: asyncio.Task = None
         self.api_helper = api_helper
-        self.lounge_controller = ytlounge.YtLoungeApi(
-            device.screen_id, config, api_helper)
         self.offset = device.offset
         self.name = device.name
         self.cancelled = False
+        self.logger = logging.getLogger(f'iSponsorBlockTV-{device.screen_id}')
+        if debug:
+            self.logger.setLevel(logging.DEBUG)
+        else:
+            self.logger.setLevel(logging.INFO)
+        rh = logging_helpers.LogHandler(device.name, log_name_len ,level=logging.DEBUG)
+        rh.add_filter_string(device.screen_id)
+        self.logger.addHandler(rh)
+        self.logger.info(f"Starting device")
+        self.lounge_controller = ytlounge.YtLoungeApi(
+            device.screen_id, config, api_helper, self.logger)
+
 
     # Ensures that we have a valid auth token
     async def refresh_auth_loop(self):
@@ -37,6 +52,7 @@ async def loop(self):
         lounge_controller = self.lounge_controller
         while not lounge_controller.linked():
             try:
+                self.logger.debug("Refreshing auth")
                 await lounge_controller.refresh_auth()
             except:
                 # traceback.print_exc()
@@ -55,18 +71,16 @@ async def loop(self):
                     await lounge_controller.connect()
                 except:
                     pass
-            print(f"Connected to device {lounge_controller.screen_name} ({self.name})")
+            self.logger.info(f"Connected to device {lounge_controller.screen_name} ({self.name})")
             try:
                 # print("Subscribing to lounge")
                 sub = await lounge_controller.subscribe_monitored(self)
                 await sub
-                await asyncio.sleep(10)
             except:
                 pass
 
     # Method called on playback state change
     async def __call__(self, state):
-        logging.debug("Playstatus update" + str(state))
         try:
             self.task.cancel()
         except:
@@ -80,7 +94,7 @@ async def process_playstatus(self, state, time_start):
         if state.videoId:
             segments = await self.api_helper.get_segments(state.videoId)
         if state.state.value == 1:  # Playing
-            print(f"Playing {state.videoId} with {len(segments)} segments")
+            self.logger.info(f"Playing video {state.videoId} with {len(segments)} segments")
             if segments:  # If there are segments
                 await self.time_to_segment(segments, state.currentTime, time_start)
 
@@ -106,10 +120,9 @@ async def time_to_segment(self, segments, position, time_start):
     # Skips to the next segment (waits for the time to pass)
     async def skip(self, time_to, position, UUID):
         await asyncio.sleep(time_to)
-        await asyncio.gather(
-            self.lounge_controller.seek_to(position),
-            self.api_helper.mark_viewed_segments(UUID)
-            )
+        self.logger.info(f"Skipping segment: seeking to {position}")
+        asyncio.create_task(self.api_helper.mark_viewed_segments(UUID))
+        asyncio.create_task(self.lounge_controller.seek_to(position))
 
     # Stops the connection to the device
     async def cancel(self):
@@ -135,11 +148,15 @@ def main(config, debug):
     tcp_connector = aiohttp.TCPConnector(ttl_dns_cache=300)
     web_session = aiohttp.ClientSession(loop=loop, connector=tcp_connector)
     api_helper = api_helpers.ApiHelper(config, web_session)
+    longest_name_len = len(list(sorted([i.name for i in config.devices]))[-1])
     for i in config.devices:
-        device = DeviceListener(api_helper, config, i)
+        device = DeviceListener(api_helper, config, i, longest_name_len, debug)
         devices.append(device)
         tasks.append(loop.create_task(device.loop()))
         tasks.append(loop.create_task(device.refresh_auth_loop()))
+
+    signal(SIGINT, lambda s, f: loop.create_task(finish(devices)))
+    rich.reconfigure(color_system="standard")
     try:
         loop.run_forever()
     except KeyboardInterrupt:

src/iSponsorBlockTV/ytlounge.py

@@ -4,23 +4,21 @@
 import pyytlounge
 from .constants import youtube_client_blacklist
 
-# Temporary imports
-from pyytlounge.api import api_base
-from pyytlounge.wrapper import NotLinkedException, desync
-
 create_task = asyncio.create_task
 
 
 class YtLoungeApi(pyytlounge.YtLoungeApi):
-    def __init__(self, screen_id, config=None, api_helper=None):
-        super().__init__("iSponsorBlockTV")
+    def __init__(self, screen_id, config=None, api_helper=None, logger=None):
+        super().__init__("iSponsorBlockTV", logger=logger)
         self.auth.screen_id = screen_id
         self.auth.lounge_id_token = None
         self.api_helper = api_helper
         self.volume_state = {}
         self.subscribe_task = None
         self.subscribe_task_watchdog = None
         self.callback = None
+        self.logger = logger
+        self.shorts_disconnected = False
         if config:
             self.mute_ads = config.mute_ads
             self.skip_ads = config.skip_ads
@@ -46,7 +44,7 @@ async def subscribe_monitored(self, callback):
 
     # Process a lounge subscription event
     def _process_event(self, event_id: int, event_type: str, args):
-        # print(f"YtLoungeApi.__process_event({event_id}, {event_type}, {args})")
+        self.logger.debug(f"process_event({event_id}, {event_type}, {args})")
         # (Re)start the watchdog
         try:
             self.subscribe_task_watchdog.cancel()
@@ -64,8 +62,6 @@ def _process_event(self, event_id: int, event_type: str, args):
                 create_task(self.mute(False, override=True))
         elif event_type == "nowPlaying":
             data = args[0]
-            self.state = pyytlounge.PlaybackState(self._logger, data)
-            self._update_state()
             # Unmute when the video starts playing
             if self.mute_ads and data.get("state", "0") == "1":
                 # print("Ad has ended, unmuting")
@@ -76,11 +72,11 @@ def _process_event(self, event_id: int, event_type: str, args):
                 # print("Ad has ended, unmuting")
                 create_task(self.mute(False, override=True))
             elif self.skip_ads and data["isSkipEnabled"] == "true":  # YouTube uses strings for booleans
-                print("Ad can be skipped, skipping")
+                self._logger.info("Ad can be skipped, skipping")
                 create_task(self.skip_ad())
                 create_task(self.mute(False, override=True))
             elif self.mute_ads:  # Seen multiple other adStates, assuming they are all ads
-                print("Ad has started, muting")
+                self._logger.info("Ad has started, muting")
                 create_task(self.mute(True, override=True))
         # Manages volume, useful since YouTube wants to know the volume when unmuting (even if they already have it)
         elif event_type == "onVolumeChanged":
@@ -98,9 +94,14 @@ def _process_event(self, event_id: int, event_type: str, args):
             data = args[0]
             # Gets segments for the next video (after the ad) before it starts playing
             if vid_id := data["contentVideoId"]:
-                print(f"Getting segments for next video: {vid_id}")
+                self._logger.info(f"Getting segments for next video: {vid_id}")
                 create_task(self.api_helper.get_segments(vid_id))
-            if self.mute_ads:
+            elif self.skip_ads and data["isSkipEnabled"] == "true":  # YouTube uses strings for booleans
+                self._logger.info("Ad can be skipped, skipping")
+                create_task(self.skip_ad())
+                create_task(self.mute(False, override=True))
+            elif self.mute_ads:  # Seen multiple other adStates, assuming they are all ads
+                self._logger.info("Ad has started, muting")
                 create_task(self.mute(True, override=True))
 
         elif event_type == "loungeStatus":
@@ -115,6 +116,17 @@ def _process_event(self, event_id: int, event_type: str, args):
         # elif event_type == "onAutoplayModeChanged":
         #     data = args[0]
         #     create_task(self.set_auto_play_mode(data["autoplayMode"] == "ENABLED"))
+        elif event_type == "onSubtitlesTrackChanged":
+            if self.shorts_disconnected:
+                data = args[0]
+                video_id_saved = data.get("videoId", None)
+                self.shorts_disconnected = False
+                create_task(self.play_video(video_id_saved))
+        elif event_type == "loungeScreenDisconnected":
+            data = args[0]
+            if data["reason"] == "disconnectedByUserScreenInitiated": # Short playing?
+                self.shorts_disconnected = True
+
         super()._process_event(event_id, event_type, args)
 
     # Set the volume to a specific value (0-100)
@@ -137,3 +149,6 @@ async def mute(self, mute: bool, override: bool = False) -> None:
 
     async def set_auto_play_mode(self, enabled: bool):
         await super()._command("setAutoplayMode", {"autoplayMode": "ENABLED" if enabled else "DISABLED"})
+
+    async def play_video(self, video_id: str) -> bool:
+        return await self._command("setPlaylist", {"videoId": video_id})