diff --git a/libs/subliminal_patch/providers/utils.py b/libs/subliminal_patch/providers/utils.py index d6fc41bdf..bf3aa856b 100644 --- a/libs/subliminal_patch/providers/utils.py +++ b/libs/subliminal_patch/providers/utils.py @@ -17,7 +17,7 @@ from ._agent_list import FIRST_THOUSAND_OR_SO_USER_AGENTS logger = logging.getLogger(__name__) -_MatchingSub = namedtuple("_MatchingSub", ("file", "priority")) +_MatchingSub = namedtuple("_MatchingSub", ("file", "priority", "context")) def _get_matching_sub( @@ -37,7 +37,7 @@ def _get_matching_sub( # If it's a movie then get the first subtitle if episode is None and episode_title is None: logger.debug("Movie subtitle found: %s", sub_name) - matching_subs.append(_MatchingSub(sub_name, 2)) + matching_subs.append(_MatchingSub(sub_name, 2, "Movie subtitle")) break guess = guessit(sub_name, options=guess_options) @@ -46,17 +46,14 @@ def _get_matching_sub( if matched_episode_num: logger.debug("No episode number found in file: %s", sub_name) - matched_title = None if episode_title is not None: - matched_title = _analize_sub_name(sub_name, episode_title) + from_name = _analize_sub_name(sub_name, episode_title) + if from_name is not None: + matching_subs.append(from_name) if episode == matched_episode_num: logger.debug("Episode matched from number: %s", sub_name) - matching_subs.append(_MatchingSub(sub_name, 2)) - elif matched_title: - matching_subs.append(_MatchingSub(sub_name, 1)) - else: - logger.debug("Ignoring incorrect episode: '%s'", sub_name) + matching_subs.append(_MatchingSub(sub_name, 2, "Episode number matched")) if matching_subs: matching_subs.sort(key=lambda x: x.priority, reverse=True) @@ -70,15 +67,21 @@ def _get_matching_sub( def _analize_sub_name(sub_name: str, title_): titles = re.split(r"[.-]", os.path.splitext(sub_name)[0]) for title in titles: + title = title.strip() ratio = SequenceMatcher(None, title, title_).ratio() if ratio > 0.85: logger.debug( "Episode title matched: '%s' -> '%s' [%s]", title, sub_name, ratio ) - return True - logger.debug("No episode title matched from file") - return False + # Avoid false positives with short titles + if len(title_) > 4 and ratio >= 0.98: + return _MatchingSub(sub_name, 3, "Perfect title ratio") + + return _MatchingSub(sub_name, 1, "Normal title ratio") + + logger.debug("No episode title matched from file: %s", sub_name) + return None def get_subtitle_from_archive( diff --git a/tests/subliminal_patch/test_utils.py b/tests/subliminal_patch/test_utils.py index 409b4eb18..eb8dc2421 100644 --- a/tests/subliminal_patch/test_utils.py +++ b/tests/subliminal_patch/test_utils.py @@ -28,26 +28,6 @@ def test_get_matching_sub(sub_names, episode, forced, expected): def test_get_matching_sub_complex_season_pack(): - files = [ - "30. Hard Drive Courage. The Ride Of The Valkyries.srt", - "34. So In Louvre Are We Two. Night Of The Scarecrow.srt", - "31. Scuba Scuba Doo. Conway The Contaminationist.srt", - "32. Katz Under The Sea. Curtain Of Cruelty.srt", - "27. Muriel Meets Her Match. Courage Vs. Mecha-Courage.srt", - "36. Fishy Business. Angry Nasty People.srt", - "28. Campsite Of Terror. The Record Deal.srt", - "33. Feast Of The Bullfrogs. Tulip's Worm.srt", - "37. Dome Of Doom. Snowman's Revenge.srt", - "35. Mondo Magic. Watch The Birdies.srt", - "29. Stormy Weather. The Sandman Sleeps.srt", - "38. The Quilt Club. Swindlin' Wind.srt", - ] - # Courage the Cowardly Dog S03E17 "Mondo Magic" - matched = utils._get_matching_sub(files, False, 17, episode_title="Mondo Magic") - assert matched == "35. Mondo Magic. Watch The Birdies.srt" - - -def test_get_matching_sub_complex_season_pack_mixed_files(): files = [ "30. Hard Drive Courage. The Ride Of The Valkyries.srt", "S03E15.srt", @@ -62,7 +42,56 @@ def test_get_matching_sub_complex_season_pack_mixed_files(): ] # Courage the Cowardly Dog S03E17 "Mondo Magic" matched = utils._get_matching_sub(files, False, 17, episode_title="Mondo Magic") - assert matched == "S03E17.srt" + assert matched == "35. Mondo Magic. Watch The Birdies.srt" + + +def _gen_subs(): + files = { + 15: "11b - Little Bigfoot.srt", + 14: "11a - Kiss Kiss, Bang Bang.srt", + 17: "10b - The Invaders.srt", + 18: "09a - The Tell Tale Tail.srt", + 1: "01 - The Thing That Wouldn't Stop It.srt", + 5: "03b - They Came from Down There.srt", + 4: "03a - Bad Day on the Moon.srt", + 8: "04a - The Friend for Life.srt", + 21: "08b - The Glazed McGuffin Affair.srt", + 13: "07b - It's Dangly Deever Time.srt", + 9: "04b - Dysfunction of the Gods.srt", + } + for episode, title in files.items(): + yield episode, title.split("-")[1].strip().rstrip(".srt"), title + + +@pytest.mark.parametrize("episode,title,file", _gen_subs()) +def test_get_matching_sub_complete_series_pack_mixed(episode, title, file): + files = [ + "11b - Little Bigfoot.srt", + "11a - Kiss Kiss, Bang Bang.srt", + "10b - The Invaders.srt", + "09a - The Tell Tale Tail.srt", + "01 - The Thing That Wouldn't Stop It.srt", + "03b - They Came from Down There.srt", + "03a - Bad Day on the Moon.srt", + "04a - The Friend for Life.srt", + "08b - The Glazed McGuffin Affair.srt", + "07b - It's Dangly Deever Time.srt", + "04b - Dysfunction of the Gods.srt", + "05b - A Glitch in Time.srt", + "12b - Sam & Max vs. the Uglions.srt", + "08a - Aaiiieee Robot.srt", + "02b - Max's Big Day.srt", + "05a - Big Trouble at the Earth's Core.srt", + "13 - The Final Episode.srt", + "02a - The Second Show Ever.srt", + "10a - Tonight We Love.srt", + "12a - Fools Die on Friday.srt", + "09b - The Trouble with Gary.srt", + "06b - We Drop at Dawn.srt", + "07a - Christmas Bloody Christmas.srt", + "06a - That Darn Gator.srt", + ] + assert utils._get_matching_sub(files, False, episode, title) == file def test_get_subtitle_from_archive_movie(data):