refactor: make trakt integration faster (#93)

custom_components/trakt_tv/apis/trakt.py

@@ -1,6 +1,5 @@
 """API for TraktTV bound to Home Assistant OAuth."""
 import logging
-import time
 from asyncio import gather, sleep
 from datetime import datetime
 from typing import Any, Dict
@@ -16,7 +15,7 @@
 from ..configuration import Configuration
 from ..const import API_HOST, DOMAIN
 from ..exception import TraktException
-from ..models.kind import BASIC_KINDS, NEXT_TO_WATCH_KINDS, TraktKind
+from ..models.kind import BASIC_KINDS, UPCOMING_KINDS, TraktKind
 from ..models.media import Medias
 from ..utils import cache_insert, cache_retrieve, deserialize_json
 
@@ -308,22 +307,41 @@ async def fetch_upcoming(
         return trakt_kind, Medias(new_medias)
 
     async def fetch_next_to_watch(
-        self, only_aired: bool = False, only_upcoming: bool = False
+        self,
+        configured_kind: TraktKind,
+        only_aired: bool = False,
+        only_upcoming: bool = False,
     ):
-        data = await gather(
-            *[
-                self.fetch_upcoming(kind, False, True, only_aired, only_upcoming)
-                for kind in NEXT_TO_WATCH_KINDS
-            ]
+        data = await self.fetch_upcoming(
+            configured_kind,
+            False,
+            True,
+            only_aired,
+            only_upcoming,
         )
-        data = filter(lambda x: x is not None, data)
-        return {trakt_kind: medias for trakt_kind, medias in data}
 
-    async def fetch_upcomings(self, all_medias: bool):
+        if data is None:
+            return {}
+
+        return dict([data])
+
+    async def fetch_upcomings(
+        self, configured_kinds: list[TraktKind], all_medias: bool
+    ):
+        kinds = []
+
+        for kind in configured_kinds:
+            if kind in UPCOMING_KINDS:
+                kinds.append(kind)
+            else:
+                LOGGER.warn(
+                    f"Upcomings doesn't support {kind}, you should remove it from the configuration."
+                )
+
         data = await gather(
             *[
                 self.fetch_upcoming(kind, all_medias, False, False, False)
-                for kind in TraktKind
+                for kind in kinds
             ]
         )
         data = filter(lambda x: x is not None, data)
@@ -334,7 +352,17 @@ async def fetch_recommendation(self, path: str, max_items: int):
             "get", f"recommendations/{path}?limit={max_items}&ignore_collected=false"
         )
 
-    async def fetch_recommendations(self):
+    async def fetch_recommendations(self, configured_kinds: list[TraktKind]):
+        kinds = []
+
+        for kind in configured_kinds:
+            if kind in BASIC_KINDS:
+                kinds.append(kind)
+            else:
+                LOGGER.warn(
+                    f"Recommendation doesn't support {kind}, you should remove it from the configuration."
+                )
+
         configuration = Configuration(data=self.hass.data)
         language = configuration.get_language()
         data = await gather(
@@ -343,11 +371,13 @@ async def fetch_recommendations(self):
                     kind.value.path,
                     configuration.get_recommendation_max_medias(kind.value.identifier),
                 )
-                for kind in BASIC_KINDS
+                for kind in kinds
             ]
         )
+
         res = {}
-        for trakt_kind, raw_medias in zip(BASIC_KINDS, data):
+
+        for trakt_kind, raw_medias in zip(kinds, data):
             if raw_medias is not None:
                 medias = [
                     trakt_kind.value.model.from_trakt(media) for media in raw_medias
@@ -356,20 +386,39 @@ async def fetch_recommendations(self):
                     *[media.get_more_information(language) for media in medias]
                 )
                 res[trakt_kind] = Medias(medias)
+
         return res
 
     async def retrieve_data(self):
         async with timeout(1800):
-            titles = []
-            data_function = []
             configuration = Configuration(data=self.hass.data)
+
+            sources = []
+            coroutine_sources_data = []
+
             source_function = {
-                "upcoming": self.fetch_upcomings(all_medias=False),
-                "all_upcoming": self.fetch_upcomings(all_medias=True),
-                "recommendation": self.fetch_recommendations(),
-                "all": self.fetch_next_to_watch(),
-                "only_aired": self.fetch_next_to_watch(only_aired=True),
-                "only_upcoming": self.fetch_next_to_watch(only_upcoming=True),
+                "upcoming": lambda kinds: self.fetch_upcomings(
+                    configured_kinds=kinds,
+                    all_medias=False,
+                ),
+                "all_upcoming": lambda kinds: self.fetch_upcomings(
+                    configured_kinds=kinds,
+                    all_medias=True,
+                ),
+                "recommendation": lambda kinds: self.fetch_recommendations(
+                    configured_kinds=kinds,
+                ),
+                "all": lambda: self.fetch_next_to_watch(
+                    configured_kind=TraktKind.NEXT_TO_WATCH_ALL,
+                ),
+                "only_aired": lambda: self.fetch_next_to_watch(
+                    configured_kind=TraktKind.NEXT_TO_WATCH_AIRED,
+                    only_aired=True,
+                ),
+                "only_upcoming": lambda: self.fetch_next_to_watch(
+                    configured_kind=TraktKind.NEXT_TO_WATCH_UPCOMING,
+                    only_upcoming=True,
+                ),
             }
 
             """First, let's configure which sensors we need depending on configuration"""
@@ -379,18 +428,23 @@ async def retrieve_data(self):
                 "recommendation",
             ]:
                 if configuration.source_exists(source):
-                    titles.append(source)
-                    data_function.append(source_function.get(source))
+                    sources.append(source)
+                    kinds = configuration.get_kinds(source)
+                    coroutine_sources_data.append(source_function.get(source)(kinds))
 
             """Then, let's add the next to watch sensors if needed"""
-            for identifier in [
+            for sub_source in [
                 "all",
                 "only_aired",
                 "only_upcoming",
             ]:
-                if configuration.next_to_watch_identifier_exists(identifier):
-                    titles.append(identifier)
-                    data_function.append(source_function.get(identifier))
+                if configuration.next_to_watch_identifier_exists(sub_source):
+                    sources.append(sub_source)
+                    coroutine_sources_data.append(source_function.get(sub_source)())
+
+            sources_data = await gather(*coroutine_sources_data)
 
-            data = await gather(*data_function)
-            return {title: medias for title, medias in zip(titles, data)}
+            return {
+                source: source_data
+                for source, source_data in zip(sources, sources_data)
+            }

custom_components/trakt_tv/configuration.py

@@ -5,6 +5,7 @@
 from dateutil.tz import tzlocal
 
 from custom_components.trakt_tv.const import DOMAIN
+from custom_components.trakt_tv.models.kind import TraktKind
 
 
 @dataclass
@@ -83,3 +84,9 @@ def source_exists(self, source: str) -> bool:
             return True
         except KeyError:
             return False
+
+    def get_kinds(self, source: str) -> list[TraktKind]:
+        return [
+            TraktKind.from_string(identifier)
+            for identifier in self.conf["sensors"][source].keys()
+        ]

custom_components/trakt_tv/models/kind.py

@@ -24,8 +24,27 @@ class TraktKind(Enum):
         "only_upcoming", "Only Upcoming", "shows", Show
     )
 
+    @classmethod
+    def from_string(cls, string):
+        try:
+            return cls[string.upper()]
+        except KeyError:
+            raise ValueError(f"No enum member found for '{string}'")
+
+
+UPCOMING_KINDS = [
+    TraktKind.SHOW,
+    TraktKind.NEW_SHOW,
+    TraktKind.PREMIERE,
+    TraktKind.MOVIE,
+    TraktKind.DVD,
+]
+
+BASIC_KINDS = [
+    TraktKind.SHOW,
+    TraktKind.MOVIE,
+]
 
-BASIC_KINDS = [TraktKind.SHOW, TraktKind.MOVIE]
 NEXT_TO_WATCH_KINDS = [
     TraktKind.NEXT_TO_WATCH_ALL,
     TraktKind.NEXT_TO_WATCH_AIRED,