feat: creation of 3 new independant next to watch sensors (#53)

README.md

@@ -84,13 +84,20 @@ trakt_tv:
         days_to_fetch: 90
         max_medias: 3
     next_to_watch:
-      show:
+      all:
         max_medias: 40
-        only_aired: True
         exclude:
           - veep
           - the-original
           - friends
+      only_aired:
+        max_medias: 40
+        exclude:
+          - veep
+          - the-original
+          - friends
+      only_upcoming:
+        max_medias: 40
 ```
 
 #### Integration Settings
@@ -158,12 +165,13 @@ Next To Watch sensor is sensor giving the next show to watch depending on your p
 
 There only one sensor available under the `sensors` > `next_to_watch` array:
 
-- `show` for [TV Shows progress](https://trakt.tv/users/<username>/progress) (actually, episodes). Creates `sensor.trakt_next_to_watch_shows`
+- `all` for all TV Shows progress. Creates `sensor.trakt_next_to_watch_all`
+- `only_aired` for only aired TV Shows progress. Creates `sensor.trakt_next_to_watch_only_aired`
+- `only_upcoming` for only upcoming TV Shows progress. Creates `sensor.trakt_next_to_watch_only_upcoming`
 
 There are three parameters for each sensor:
 
 - `max_medias` should be a positive number for how many items to grab
-- `only_aired` should be a boolean (True / False) to display only already aired episodes
 - `exclude` should be a list of shows you'd like to exclude, since it's based on your watched history. To find keys to put there, go on trakt.tv, search for a show, click on it, notice the url slug, copy/paste it. So, if I want to hide "Friends", I'll do the steps mentioned above, then land on https://trakt.tv/shows/friends, I'll just have to copy/paste the last part, `friends`, that's it
 You can also use the Trakt.tv "hidden" function to hide a show from [your calendar](https://trakt.tv/calendars/my/shows) or the [progress page](https://trakt.tv/users/<username>/progress)
 

custom_components/trakt_tv/apis/trakt.py

@@ -14,7 +14,7 @@
 
 from ..configuration import Configuration
 from ..const import API_HOST, DOMAIN
-from ..models.kind import BASIC_KINDS, TraktKind
+from ..models.kind import BASIC_KINDS, NEXT_TO_WATCH_KINDS, TraktKind
 from ..models.media import Medias
 from ..utils import deserialize_json
 
@@ -141,7 +141,12 @@ async def fetch_show_informations(
         )
 
     async def fetch_upcoming(
-        self, trakt_kind: TraktKind, all_medias: bool, next_to_watch: bool
+        self,
+        trakt_kind: TraktKind,
+        all_medias: bool,
+        next_to_watch: bool,
+        only_aired: bool,
+        only_upcoming: bool,
     ):
         """
         Fetch the calendar of the user trakt account based on the trak_type containing
@@ -156,23 +161,25 @@ async def fetch_upcoming(
         path = trakt_kind.value.path
         identifier = trakt_kind.value.identifier
 
-        if (
-            not configuration.upcoming_identifier_exists(identifier, all_medias)
-        ) and not (configuration.next_to_watch_identifier_exists(identifier)):
-            return None
+        upcoming_identifier_exists = configuration.upcoming_identifier_exists(
+            identifier, all_medias
+        )
+        next_to_watch_identifier_exists = configuration.next_to_watch_identifier_exists(
+            identifier
+        )
 
-        if next_to_watch and trakt_kind.value.identifier != "show":
+        if (next_to_watch and (not next_to_watch_identifier_exists)) or (
+            (not next_to_watch) and (not upcoming_identifier_exists)
+        ):
             return None
 
         configuration = Configuration(data=self.hass.data)
 
         max_medias = configuration.get_upcoming_max_medias(identifier, all_medias)
         language = configuration.get_language()
-        only_aired = False
 
         if next_to_watch:
             excluded_shows = configuration.get_exclude_shows(identifier)
-            only_aired = configuration.get_only_aired(identifier)
             raw_medias = await self.fetch_watched(excluded_shows)
         else:
             days_to_fetch = configuration.get_upcoming_days_to_fetch(
@@ -197,7 +204,14 @@ async def fetch_upcoming(
                     pytz.timezone(configuration.get_timezone())
                 )
                 new_medias = [
-                    media for media in medias if media.released < timezoned_now
+                    media for media in medias if media.released <= timezoned_now
+                ]
+            elif only_upcoming:
+                timezoned_now = datetime.now(
+                    pytz.timezone(configuration.get_timezone())
+                )
+                new_medias = [
+                    media for media in medias if media.released > timezoned_now
                 ]
             else:
                 new_medias = medias
@@ -209,16 +223,24 @@ async def fetch_upcoming(
 
         return trakt_kind, Medias(new_medias)
 
-    async def fetch_next_to_watch(self):
+    async def fetch_next_to_watch(
+        self, only_aired: bool = False, only_upcoming: bool = False
+    ):
         data = await gather(
-            *[self.fetch_upcoming(kind, False, True) for kind in TraktKind]
+            *[
+                self.fetch_upcoming(kind, False, True, only_aired, only_upcoming)
+                for kind in NEXT_TO_WATCH_KINDS
+            ]
         )
         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):
         data = await gather(
-            *[self.fetch_upcoming(kind, all_medias, False) for kind in TraktKind]
+            *[
+                self.fetch_upcoming(kind, all_medias, False, False, False)
+                for kind in TraktKind
+            ]
         )
         data = filter(lambda x: x is not None, data)
         return {trakt_kind: medias for trakt_kind, medias in data}
@@ -251,13 +273,22 @@ async def fetch_recommendations(self):
 
     async def retrieve_data(self):
         async with timeout(60):
-            titles = ["upcoming", "all_upcoming", "recommendation", "next_to_watch"]
+            titles = [
+                "upcoming",
+                "all_upcoming",
+                "recommendation",
+                "all",
+                "only_aired",
+                "only_upcoming",
+            ]
             data = await gather(
                 *[
                     self.fetch_upcomings(all_medias=False),
                     self.fetch_upcomings(all_medias=True),
                     self.fetch_recommendations(),
                     self.fetch_next_to_watch(),
+                    self.fetch_next_to_watch(only_aired=True),
+                    self.fetch_next_to_watch(only_upcoming=True),
                 ]
             )
             return {title: medias for title, medias in zip(titles, data)}

custom_components/trakt_tv/configuration.py

@@ -52,12 +52,6 @@ def get_exclude_shows(self, identifier: str) -> list:
         except KeyError:
             return []
 
-    def get_only_aired(self, identifier: str) -> bool:
-        try:
-            return self.conf["sensors"]["next_to_watch"][identifier]["only_aired"]
-        except KeyError:
-            return False
-
     def next_to_watch_identifier_exists(self, identifier: str) -> bool:
         return self.identifier_exists(identifier, "next_to_watch")
 

custom_components/trakt_tv/models/kind.py

@@ -18,6 +18,16 @@ class TraktKind(Enum):
     PREMIERE = CalendarInformation("premiere", "Premieres", "shows/premieres", Show)
     MOVIE = CalendarInformation("movie", "Movies", "movies", Movie)
     DVD = CalendarInformation("dvd", "DVD", "dvd", Movie)
+    NEXT_TO_WATCH_ALL = CalendarInformation("all", "All", "shows", Show)
+    NEXT_TO_WATCH_AIRED = CalendarInformation("only_aired", "Only Aired", "shows", Show)
+    NEXT_TO_WATCH_UPCOMING = CalendarInformation(
+        "only_upcoming", "Only Upcoming", "shows", Show
+    )
 
 
 BASIC_KINDS = [TraktKind.SHOW, TraktKind.MOVIE]
+NEXT_TO_WATCH_KINDS = [
+    TraktKind.NEXT_TO_WATCH_ALL,
+    TraktKind.NEXT_TO_WATCH_AIRED,
+    TraktKind.NEXT_TO_WATCH_UPCOMING,
+]

custom_components/trakt_tv/schema.py

@@ -7,7 +7,7 @@
 from voluptuous import ALLOW_EXTRA, PREVENT_EXTRA, In, Required, Schema
 
 from .const import DOMAIN, LANGUAGE_CODES
-from .models.kind import BASIC_KINDS, TraktKind
+from .models.kind import BASIC_KINDS, NEXT_TO_WATCH_KINDS, TraktKind
 
 
 def dictionary_to_schema(
@@ -57,10 +57,9 @@ def upcoming_schema() -> Dict[str, Any]:
 
 def next_to_watch_schema() -> Dict[str, Any]:
     subschemas = {}
-    for trakt_kind in TraktKind:
+    for trakt_kind in NEXT_TO_WATCH_KINDS:
         subschemas[trakt_kind.value.identifier] = {
             Required("max_medias", default=3): cv.positive_int,
-            Required("only_aired", default=False): bool,
             Required("exclude", default=[]): list,
         }
 

custom_components/trakt_tv/sensor.py

@@ -6,7 +6,7 @@
 
 from .configuration import Configuration
 from .const import DOMAIN
-from .models.kind import BASIC_KINDS, TraktKind
+from .models.kind import BASIC_KINDS, NEXT_TO_WATCH_KINDS, TraktKind
 
 LOGGER = logging.getLogger(__name__)
 
@@ -50,16 +50,19 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
             )
             sensors.append(sensor)
 
-        if trakt_kind.value.identifier != "show":
+    for trakt_kind in TraktKind:
+        if trakt_kind not in NEXT_TO_WATCH_KINDS:
             continue
 
+        identifier = trakt_kind.value.identifier
+
         if configuration.next_to_watch_identifier_exists(identifier):
             sensor = TraktSensor(
                 hass=hass,
                 config_entry=config_entry,
                 coordinator=coordinator,
                 trakt_kind=trakt_kind,
-                source="next_to_watch",
+                source=identifier,
                 prefix="Trakt Next To Watch",
                 mdi_icon="mdi:calendar",
             )
@@ -105,7 +108,10 @@ def medias(self):
     def configuration(self):
         identifier = self.trakt_kind.value.identifier
         data = self.hass.data[DOMAIN]
-        return data["configuration"]["sensors"][self.source][identifier]
+        source = (
+            "next_to_watch" if self.trakt_kind in NEXT_TO_WATCH_KINDS else self.source
+        )
+        return data["configuration"]["sensors"][source][identifier]
 
     @property
     def data(self):