diff --git a/requirements.txt b/requirements.txt index b3cf464..55f5d38 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ deprecated~=1.2.13 requests-oauthlib>=1.3 requests>=2.25 +dataclasses; python_version<"3.7" diff --git a/tests/mock_data/users.json b/tests/mock_data/users.json index 3e9176e..b555a13 100644 --- a/tests/mock_data/users.json +++ b/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}} } diff --git a/trakt/users.py b/trakt/users.py index a9412a6..2c8336d 100644 --- a/trakt/users.py +++ b/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