diff --git a/tests/mixins/test_playlists.py b/tests/mixins/test_playlists.py index 4a1ebfa1..a94fbbb2 100644 --- a/tests/mixins/test_playlists.py +++ b/tests/mixins/test_playlists.py @@ -104,6 +104,16 @@ def test_get_large_audio_playlist(self, yt_oauth): assert len(album["tracks"]) == 456 assert album["trackCount"] == 456 + def test_get_playlist_validate_responses(self, yt_oauth): + playlist_id = "PL6bPxvf5dW5clc3y9wAoslzqUrmkZ5c-u" # very large + playlist = yt_oauth.get_playlist(playlist_id, limit=None, validate_responses=True) + # validated continuations should retrieve all tracks of the playlist + assert len(playlist["tracks"]) == playlist["trackCount"] + + def test_get_playlist_validate_responses_no_track_count(self, yt_oauth): + with pytest.raises(YTMusicUserError): + yt_oauth.get_playlist("RDATgXd-", limit=None, validate_responses=True) + @pytest.mark.parametrize( "playlist_id", [ diff --git a/tests/test_continuations.py b/tests/test_continuations.py new file mode 100644 index 00000000..96c69b1a --- /dev/null +++ b/tests/test_continuations.py @@ -0,0 +1,98 @@ +from ytmusicapi.continuations import get_validated_continuations_2025 +from ytmusicapi.type_alias import JsonDict, JsonList + + +def _make_page(video_ids: list[int], next_token: str | None) -> JsonDict: + """Build a fake YouTube Music continuation response.""" + continuation_items: JsonList = [{"videoId": video_id} for video_id in video_ids] + if next_token is not None: + continuation_items.append( + { + "continuationItemRenderer": { + "continuationEndpoint": {"continuationCommand": {"token": next_token}} + } + } + ) + return { + "onResponseReceivedActions": [ + {"appendContinuationItemsAction": {"continuationItems": continuation_items}} + ] + } + + +def _parse_func(continuation_items: JsonList) -> JsonList: + return [item for item in continuation_items if "videoId" in item] + + +def _initial_results(first_token: str) -> JsonDict: + return { + "contents": [ + {"videoId": 0}, + { + "continuationItemRenderer": { + "continuationEndpoint": {"continuationCommand": {"token": first_token}} + } + }, + ] + } + + +class TestValidatedContinuations2025: + def test_retries_short_pages(self) -> None: + """A continuation page that returns too few items is retried until it is complete.""" + t1_calls = 0 + + def request_func(body: JsonDict) -> JsonDict: + nonlocal t1_calls + token = body["continuation"] + if token == "T1": + t1_calls += 1 + if t1_calls == 1: + return _make_page([1], "T2") # short response, must be retried + return _make_page([1, 2, 3], "T2") + if token == "T2": + return _make_page([4, 5, 6], None) # last page, no further continuation + raise AssertionError(f"unexpected token {token}") + + items = get_validated_continuations_2025( + _initial_results("T1"), limit=6, per_page=3, request_func=request_func, parse_func=_parse_func + ) + + assert [item["videoId"] for item in items] == [1, 2, 3, 4, 5, 6] + assert t1_calls == 2 # first attempt was short and retried once + + def test_gives_up_after_max_retries(self) -> None: + """If a page never returns enough items, retries are capped and partial results are returned.""" + calls = 0 + + def request_func(body: JsonDict) -> JsonDict: + nonlocal calls + calls += 1 + return _make_page([1], None) # always short, no next token + + items = get_validated_continuations_2025( + _initial_results("T1"), + limit=3, + per_page=3, + request_func=request_func, + parse_func=_parse_func, + max_retries=2, + ) + + assert [item["videoId"] for item in items] == [1] + assert calls == 3 # 1 initial attempt + 2 retries + + def test_stops_at_limit(self) -> None: + """No continuation request is sent once the limit is reached.""" + + def request_func(body: JsonDict) -> JsonDict: + token = body["continuation"] + if token == "T1": + return _make_page([1, 2, 3], "T2") + raise AssertionError("continuation past the limit should not be requested") + + items = get_validated_continuations_2025( + _initial_results("T1"), limit=3, per_page=3, request_func=request_func, parse_func=_parse_func + ) + + assert [item["videoId"] for item in items] == [1, 2, 3] diff --git a/ytmusicapi/continuations.py b/ytmusicapi/continuations.py index 43fa948f..6ecad66c 100644 --- a/ytmusicapi/continuations.py +++ b/ytmusicapi/continuations.py @@ -1,5 +1,5 @@ from collections.abc import Callable -from typing import Any, cast +from typing import Any, TypeVar, cast from ytmusicapi.navigation import nav from ytmusicapi.type_alias import ( @@ -20,6 +20,8 @@ ] CONTINUATION_ITEMS = ["onResponseReceivedActions", 0, "appendContinuationItemsAction", "continuationItems"] +T = TypeVar("T") + def get_continuation_token(results: JsonList) -> str | None: last_result = results[-1] @@ -60,6 +62,56 @@ def get_continuations_2025( return items +def get_validated_continuations_2025( + results: JsonDict, + limit: int, + per_page: int, + request_func: RequestFuncBodyType, + parse_func: ParseFuncType, + max_retries: int = 3, +) -> JsonList: + """Validated variant of :py:func:`get_continuations_2025`. + + Some YouTube Music responses return fewer items than expected for a continuation. + This retries each continuation request until the expected amount of items is returned. + + :param results: result dict containing the initial ``contents`` and continuation token + :param limit: minimum number of items to retrieve in total. Whole pages are fetched, so the + result may exceed this value + :param per_page: expected number of items per continuation request + :param request_func: request func that accepts a body dict (see :py:data:`RequestFuncBodyType`) + :param parse_func: parse func to apply on the returned continuation items + :param max_retries: how often to retry a single continuation request before giving up + :return: list of parsed continuation results + """ + items: JsonList = [] + continuation_token = get_continuation_token(results["contents"]) + while continuation_token and len(items) < limit: + wrapped_parse_func = lambda raw_response: get_parsed_continuation_items_2025(raw_response, parse_func) + validate_func = lambda parsed: validate_response(parsed, per_page, limit, len(items)) + + response = resend_request_until_parsed_response_is_valid( + request_func, {"continuation": continuation_token}, wrapped_parse_func, validate_func, max_retries + ) + continuation_items = response["results"] + if not continuation_items: + break + + parsed = response["parsed"] + if len(parsed) == 0: + break + items.extend(parsed) + continuation_token = get_continuation_token(continuation_items) + + return items + + +def get_parsed_continuation_items_2025(response: JsonDict, parse_func: ParseFuncType) -> JsonDict: + continuation_items = nav(response, CONTINUATION_ITEMS, True) or [] + parsed = parse_func(continuation_items) if continuation_items else [] + return {"results": continuation_items, "parsed": parsed} + + def get_reloadable_continuations( results: JsonDict, continuation_type: str, @@ -174,8 +226,8 @@ def get_continuation_contents(continuation: JsonDict, parse_func: ParseFuncType) def resend_request_until_parsed_response_is_valid( - request_func: RequestFuncType, - request_additional_params: str, + request_func: Callable[[T], JsonDict], + request_additional_params: T, parse_func: ParseFuncDictType, validate_func: Callable[[dict[str, Any]], bool], max_retries: int, diff --git a/ytmusicapi/mixins/playlists.py b/ytmusicapi/mixins/playlists.py index bdc2476e..01e1b297 100644 --- a/ytmusicapi/mixins/playlists.py +++ b/ytmusicapi/mixins/playlists.py @@ -16,7 +16,12 @@ class PlaylistsMixin(MixinProtocol): def get_playlist( - self, playlistId: str, limit: int | None = 100, related: bool = False, suggestions_limit: int = 0 + self, + playlistId: str, + limit: int | None = 100, + related: bool = False, + suggestions_limit: int = 0, + validate_responses: bool = False, ) -> JsonDict: """ Returns a list of playlist items @@ -27,6 +32,9 @@ def get_playlist( :param suggestions_limit: How many suggestions to return. The result is a list of suggested playlist items (videos) contained in a "suggestions" key. 7 items are retrieved in each internal request. Default: 0 + :param validate_responses: Flag indicating if responses from YTM should be validated and retried + in case when some tracks are missing. This helps to retrieve all tracks of large playlists, + where YouTube Music sometimes returns fewer items than expected per continuation. Default: False :return: Dictionary with information about the playlist. The key ``tracks`` contains a List of playlistItem dictionaries @@ -225,9 +233,33 @@ def get_playlist( ) parse_func = lambda contents: parse_playlist_items(contents, is_collaborative=is_collaborative) - playlist["tracks"].extend( - get_continuations_2025(content_data, limit, request_func_continuations, parse_func) - ) + if validate_responses: + per_page = 100 + track_count = playlist["trackCount"] + if limit is None: + request_count = track_count + elif track_count is None: + request_count = limit + else: + request_count = min(limit, track_count) + if request_count is None: + raise YTMusicUserError( + "Validation is not supported for playlists without a known trackCount. " + "Please provide a limit." + ) + playlist["tracks"].extend( + get_validated_continuations_2025( + content_data, + request_count - len(playlist["tracks"]), + per_page, + request_func_continuations, + parse_func, + ) + ) + else: + playlist["tracks"].extend( + get_continuations_2025(content_data, limit, request_func_continuations, parse_func) + ) playlist["duration_seconds"] = sum_total_duration(playlist) return playlist