diff --git a/tests/core/test_load_tracks_enrich.py b/tests/core/test_load_tracks_enrich.py new file mode 100644 index 00000000..fabe9ca0 --- /dev/null +++ b/tests/core/test_load_tracks_enrich.py @@ -0,0 +1,44 @@ +from types import SimpleNamespace + +from langcodes import Language + +from unshackle.commands.dl import apply_enrich, load_tracks +from unshackle.core.remote_service import apply_title_fields +from unshackle.core.titles import Episode, Movie, Song +from unshackle.core.tracks import Chapters, Tracks + +SERVICE = type("SVC", (), {}) + + +def refreshing_service(info): + """A service that sets its title again in get_tracks, as a remote session does.""" + return SimpleNamespace( + get_tracks=lambda title: apply_title_fields(title, info) or Tracks(), + get_chapters=lambda title: Chapters(), + ) + + +def test_enrich_wins_over_a_movie_title_set_in_get_tracks(): + movie = Movie(id_="movie-0001", service=SERVICE, name="The Gold Seekers", year=2017) + service = refreshing_service({"name": "Los Buscadores", "year": 2016, "language": "pt"}) + load_tracks(service, movie, name="The Gold Seekers", year=2017, language=Language.get("es")) + assert (movie.name, movie.year, str(movie.language)) == ("The Gold Seekers", 2017, "es") + + +def test_enrich_wins_over_a_series_title_set_in_get_tracks(): + episode = Episode(id_="episode-0001", service=SERVICE, title="English", season=1, number=1, name="Pilot") + service = refreshing_service({"series_title": "Local", "name": "Pilot", "year": 2016}) + load_tracks(service, episode, name="English", year=2017) + assert (episode.title, episode.name, episode.year) == ("English", "Pilot", 2017) + + +def test_no_enrich_keeps_the_title_from_get_tracks(): + movie = Movie(id_="movie-0001", service=SERVICE, name="Old", year=2016) + load_tracks(refreshing_service({"name": "New", "year": 2017}), movie) + assert (movie.name, movie.year) == ("New", 2017) + + +def test_a_song_takes_only_the_language(): + song = Song(id_="song-0001", service=SERVICE, name="Track 1", artist="A", album="B", track=1, disc=1, year=1999) + apply_enrich(song, "The Gold Seekers", 2017, Language.get("es")) + assert (song.name, song.year, str(song.language)) == ("Track 1", 1999, "es") diff --git a/unshackle/commands/dl.py b/unshackle/commands/dl.py index e64c481d..310649e1 100644 --- a/unshackle/commands/dl.py +++ b/unshackle/commands/dl.py @@ -227,6 +227,35 @@ def latest_episode_keys(titles: Any, latest_episode: bool, latest_episodes: Opti return list(dict.fromkeys((ep.season, ep.number) for ep in titles))[-count:] +def apply_enrich(title: Title_T, name: Optional[str], year: Optional[int], language: Optional[Language]) -> None: + """Set the --enrich name, year and original language on a title. A song takes only the language. + + A service can set the title again in get_tracks, so dl applies them after that call too. + """ + if isinstance(title, Episode): + title.title = name or title.title + elif isinstance(title, Movie): + title.name = name or title.name + if isinstance(title, (Episode, Movie)): + title.year = year or title.year + if language: + apply_original_language(title, language) + + +def load_tracks( + service: Service, + title: Title_T, + *, + name: Optional[str] = None, + year: Optional[int] = None, + language: Optional[Language] = None, +) -> None: + """Load the tracks and chapters of a title, then apply the --enrich name, year and language again.""" + title.tracks.add(service.get_tracks(title), warn_only=True) + title.tracks.chapters = service.get_chapters(title) + apply_enrich(title, name, year, language) + + def server_url(server_name: Optional[str]) -> str: """The configured URL of a remote server, or an empty string when the config cannot be read.""" from unshackle.core.remote_service import resolve_server @@ -1754,13 +1783,12 @@ class dl: cache_account_hash = get_account_hash(service.credential) if hasattr(service, "credential") else None enrich_lang: Optional[Language] = None + enrich_title: Optional[str] = None + enrich_year: Optional[int] = None if self.enrich: sample_title = titles[0] if hasattr(titles, "__getitem__") else titles kind = "tv" if isinstance(sample_title, Episode) else "movie" - enrich_title: Optional[str] = None - enrich_year: Optional[int] = None - enrich_result = providers.resolve_by_ids( self.tmdb_id, self.imdb_id, @@ -1789,17 +1817,9 @@ class dl: ]: self.log.warning(f"--enrich source gave no {', '.join(missing)}. Those fields are unchanged.") - if enrich_title or enrich_year or enrich_lang: - for t in titles if isinstance(titles, (Series, Movies)) else [titles]: - if enrich_title: - if isinstance(t, Episode): - t.title = enrich_title - else: - t.name = enrich_title - if enrich_year: - t.year = enrich_year - if enrich_lang: - t.language = enrich_lang + if isinstance(titles, (Series, Movies)): + for t in titles: + apply_enrich(t, enrich_title, enrich_year, enrich_lang) if isinstance(titles, Series): enrich_tvdb_id = self.tvdb_id or (enrich_result.external_ids.tvdb_id if enrich_result else None) @@ -2090,10 +2110,7 @@ class dl: tracks_label = "Getting Remote Tracks..." if self.is_remote else "Getting Tracks..." with console.status(tracks_label, spinner="dots"): try: - title.tracks.add(service.get_tracks(title), warn_only=True) - title.tracks.chapters = service.get_chapters(title) - if enrich_lang: - apply_original_language(title, enrich_lang) + load_tracks(service, title, name=enrich_title, year=enrich_year, language=enrich_lang) except Exception as e: if self.debug_logger: self.debug_logger.log_error(