From 17a23c1f7838237c43e0d83552f12889f22182dd Mon Sep 17 00:00:00 2001 From: Dylan Do Amaral <37652599+dylandoamaral@users.noreply.github.com> Date: Fri, 12 Jan 2024 19:52:14 +0100 Subject: [PATCH] refactor: make trakt integration faster (#93) --- custom_components/trakt_tv/apis/trakt.py | 116 ++++++++++++++------ custom_components/trakt_tv/configuration.py | 7 ++ custom_components/trakt_tv/models/kind.py | 21 +++- 3 files changed, 112 insertions(+), 32 deletions(-) diff --git a/custom_components/trakt_tv/apis/trakt.py b/custom_components/trakt_tv/apis/trakt.py index 776bd85..8a6a3d1 100644 --- a/custom_components/trakt_tv/apis/trakt.py +++ b/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) + } diff --git a/custom_components/trakt_tv/configuration.py b/custom_components/trakt_tv/configuration.py index 4b40513..72a897b 100644 --- a/custom_components/trakt_tv/configuration.py +++ b/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() + ] diff --git a/custom_components/trakt_tv/models/kind.py b/custom_components/trakt_tv/models/kind.py index b43ad46..b098dba 100644 --- a/custom_components/trakt_tv/models/kind.py +++ b/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,