diff --git a/tests/test_pagination.py b/tests/test_pagination.py new file mode 100644 index 0000000..864c9fa --- /dev/null +++ b/tests/test_pagination.py @@ -0,0 +1,73 @@ +from trakt.pagination import paginate + + +class FakeClient: + def __init__(self, pages): + self.pages = list(pages) + self.urls = [] + + def get(self, url, include_headers=False): + self.urls.append((url, include_headers)) + return self.pages.pop(0) + + +def test_paginate_fetches_all_list_pages(): + client = FakeClient([ + ([{'title': 'One'}], {'X-Pagination-Page-Count': '3'}), + ([{'title': 'Two'}], {'X-Pagination-Page-Count': '3'}), + ([{'title': 'Three'}], {'X-Pagination-Page-Count': '3'}), + ]) + + result = paginate( + 'users/{user}/watched/movies', + api=client, + user='sean', + limit=2, + ) + + assert result == [ + {'title': 'One'}, + {'title': 'Two'}, + {'title': 'Three'}, + ] + assert client.urls == [ + ('users/sean/watched/movies?limit=2', True), + ('users/sean/watched/movies?page=2&limit=2', True), + ('users/sean/watched/movies?page=3&limit=2', True), + ] + + +def test_paginate_stops_after_one_page_without_valid_page_count(): + client = FakeClient([ + ([{'title': 'One'}], {'X-Pagination-Page-Count': 'invalid'}), + ([{'title': 'Two'}], {'X-Pagination-Page-Count': '2'}), + ]) + + result = paginate('movies/popular', api=client) + + assert result == [{'title': 'One'}] + assert client.urls == [('movies/popular', True)] + + +def test_paginate_stops_after_one_page_without_page_count(): + client = FakeClient([ + ([{'title': 'One'}], {}), + ([{'title': 'Two'}], {'X-Pagination-Page-Count': '2'}), + ]) + + result = paginate('movies/popular', api=client) + + assert result == [{'title': 'One'}] + assert client.urls == [('movies/popular', True)] + + +def test_paginate_skips_none_extends_lists_and_appends_objects(): + client = FakeClient([ + (None, {'X-Pagination-Page-Count': '3'}), + ([{'title': 'Two'}], {'X-Pagination-Page-Count': '3'}), + ({'title': 'Three'}, {'X-Pagination-Page-Count': '3'}), + ]) + + result = paginate('movies/popular', api=client) + + assert result == [{'title': 'Two'}, {'title': 'Three'}] \ No newline at end of file diff --git a/trakt/api.py b/trakt/api.py index 4daf2e8..5ed2f56 100644 --- a/trakt/api.py +++ b/trakt/api.py @@ -12,7 +12,6 @@ from trakt.core import TIMEOUT from trakt.errors import (BadRequestException, BadResponseException, OAuthException, OAuthRefreshException) -from trakt.utils import build_uri __author__ = 'Elan Ruusamäe' @@ -80,35 +79,6 @@ def put(self, url: str, data): """ return self.request('put', url, data=data) - def iter_pages(self, url, **params): - """Yield successive pages for a paginated GET endpoint.""" - page = 1 - while True: - page_url = build_uri(url, **params) if page == 1 else build_uri( - url, page=page, **params - ) - page_data, headers = self.get(page_url, include_headers=True) - yield page_data - try: - page_count = int(headers.get('X-Pagination-Page-Count', 1)) - except (TypeError, ValueError): - page_count = 1 - if page >= page_count: - break - page += 1 - - def paginate(self, url, **params): - """Return a flattened list from all pages of a paginated GET endpoint.""" - results = [] - for page_data in self.iter_pages(url, **params): - if page_data is None: - continue - if isinstance(page_data, list): - results.extend(page_data) - else: - results.append(page_data) - return results - @property def auth(self): """ diff --git a/trakt/pagination.py b/trakt/pagination.py new file mode 100644 index 0000000..77a728b --- /dev/null +++ b/trakt/pagination.py @@ -0,0 +1,44 @@ +"""Pagination helpers for Trakt API clients.""" + +from trakt.core import api as factory +from trakt.utils import build_uri + +__author__ = 'Elan Ruusamäe' +__all__ = ['paginate'] + + +def page_count(headers): + try: + return int(headers.get('X-Pagination-Page-Count', 1)) + except (TypeError, ValueError): + return 1 + + +def iter_pages(client, url: str, **params): + """Yield successive pages for a paginated GET endpoint.""" + params.pop('page', None) + page = 1 + while True: + page_url = build_uri(url, **params) if page == 1 else build_uri( + url, page=page, **params + ) + page_data, headers = client.get(page_url, include_headers=True) + yield page_data + + if page >= page_count(headers): + break + page += 1 + + +def paginate(url: str, api=None, **params): + """Return a flattened list from all pages of a paginated GET endpoint.""" + results = [] + client = api or factory() + for page_data in iter_pages(client, url, **params): + if page_data is None: + continue + if isinstance(page_data, list): + results.extend(page_data) + else: + results.append(page_data) + return results diff --git a/trakt/users.py b/trakt/users.py index 56dd148..e72407c 100644 --- a/trakt/users.py +++ b/trakt/users.py @@ -4,9 +4,10 @@ from dataclasses import dataclass, fields from typing import Any, NamedTuple, Optional, Union -from trakt.core import api, delete, get, post +from trakt.core import delete, get, post from trakt.mixins import DataClassMixin, IdsMixin from trakt.movies import Movie +from trakt.pagination import paginate from trakt.people import Person from trakt.tv import TVEpisode, TVSeason, TVShow from trakt.utils import build_uri, slugify @@ -548,8 +549,7 @@ def watched_movies(self): collection. Automatically fetches all pages. """ if self._watched_movies is None: - client = api() - all_movies = client.paginate( + all_movies = paginate( 'users/{user}/watched/movies', user=slugify(self.username), )