Add dataclasses for Watchlist responses

requirements.txt

@@ -1,3 +1,4 @@
 deprecated~=1.2.13
 requests-oauthlib>=1.3
 requests>=2.25
+dataclasses; python_version<"3.7"

tests/mock_data/users.json

@@ -527,11 +527,17 @@
         "GET": [
             {
                 "listed_at":"2014-09-01T09:10:11.000Z",
+                "rank": 1,
+                "id": 1,
+                "notes": null,
                 "type":"movie",
                 "movie":{"title":"TRON: Legacy","year":2010,"ids":{"trakt":1,"slug":"tron-legacy-2010","imdb":"tt1104001","tmdb":20526}}
             },
             {
                 "listed_at":"2014-09-01T09:10:11.000Z",
+                "rank": 2,
+                "id": 6,
+                "notes": null,
                 "type":"movie",
                 "movie":{"title":"The Dark Knight","year":2008,"ids":{"trakt":6,"slug":"the-dark-knight-2008","imdb":"tt0468569","tmdb":155}}
             }
@@ -541,11 +547,17 @@
         "GET": [
             {
                 "listed_at":"2014-09-01T09:10:11.000Z",
+                "rank": 1,
+                "id": 1,
+                "notes": null,
                 "type":"show",
                 "show":{"title":"Breaking Bad","year":2008,"ids":{"trakt":1,"slug":"breaking-bad","tvdb":81189,"imdb":"tt0903747","tmdb":1396,"tvrage":18164}}
             },
             {
                 "listed_at":"2014-09-01T09:10:11.000Z",
+                "rank": 2,
+                "id": 2,
+                "notes": null,
                 "type":"show",
                 "show":{"title":"The Walking Dead","year":2010,"ids":{"trakt":2,"slug":"the-walking-dead","tvdb":153021,"imdb":"tt1520211","tmdb":1402,"tvrage":null}}
             }

trakt/users.py

@@ -1,6 +1,9 @@
 # -*- coding: utf-8 -*-
 """Interfaces to all of the User objects offered by the Trakt.tv API"""
+
 from collections import namedtuple
+from dataclasses import dataclass
+from typing import Any, Optional
 
 from trakt.core import delete, get, post
 from trakt.mixins import IdsMixin
@@ -9,8 +12,11 @@
 from trakt.tv import TVEpisode, TVSeason, TVShow
 from trakt.utils import slugify
 
-__author__ = 'Jon Nappi'
-__all__ = ['User', 'UserList', 'Request', 'follow', 'get_all_requests',
+__author__ = 'Jon Nappi, Elan Ruusamäe'
+__all__ = ['User', 'UserList',
+           'WatchlistMovieEntry',
+           'WatchlistShowEntry',
+           'Request', 'follow', 'get_all_requests',
            'get_user_settings', 'unfollow']
 
 
@@ -26,6 +32,31 @@ def deny(self):
         yield 'users/requests/{id}'.format(id=self.id)
 
 
+@dataclass
+class WatchlistEntry:
+    rank: int
+    id: int
+    listed_at: str
+    notes: Optional[str]
+    type: Literal["movies", "shows"]
+
+
+@dataclass
+class WatchlistMovieEntry(Movie, WatchlistEntry):
+    movie: Any
+
+    def __post_init__(self):
+        super().__init__(**self.movie)
+
+
+@dataclass
+class WatchlistShowEntry(TVShow, WatchlistEntry):
+    show: Any
+
+    def __post_init__(self):
+        super().__init__(**self.show)
+
+
 @post
 def follow(user_name):
     """Follow a user with *user_name*. If the user has a protected profile, the
@@ -320,9 +351,8 @@ def watchlist_shows(self):
             )
             self._show_watchlist = []
             for show in data:
-                show_data = show.pop('show')
-                show_data.update(show)
-                self._show_watchlist.append(TVShow(**show_data))
+                wl = WatchlistShowEntry(**show)
+                self._show_watchlist.append(wl)
             yield self._show_watchlist
         yield self._show_watchlist
 
@@ -337,8 +367,8 @@ def watchlist_movies(self):
             )
             self._movie_watchlist = []
             for movie in data:
-                mov = movie.pop('movie')
-                self._movie_watchlist.append(Movie(**mov))
+                wl = WatchlistMovieEntry(**movie)
+                self._movie_watchlist.append(wl)
             yield self._movie_watchlist
         yield self._movie_watchlist