Files
HS-Rename/engine/episode_match.py
T
Bulk RenamerandCursor 85dc873e16 Release v1.0.8: recursive folder scan and file type filters.
Adds subfolder inclusion and extension presets so the preview list can target nested media without loading everything.

Co-authored-by: Cursor <[email protected]>
2026-07-03 17:46:58 -05:00

261 lines
7.9 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
Match episode titles from filenames against a reference episode list (e.g. TheTVDB).
"""
from __future__ import annotations
import re
from dataclasses import dataclass
from difflib import SequenceMatcher
from pathlib import Path
from typing import Optional
from .tvdb_client import TvdbEpisode
# S01E05 - Title or S01E05-E06 - Title
PATTERN_SXXEXX = re.compile(
r"^(.*?[Ss])(\d+)([Ee])(\d+)(-[Ee](\d+))?(.*)$",
)
# Show Name 04x01 Title, 4x01 - Title, etc.
PATTERN_NXNN = re.compile(
r"^(.*?)(\d{1,2})[xX](\d{1,4})(?:([\s._-]+)(.+))?$",
)
DEFAULT_EPISODE_PATTERN = PATTERN_SXXEXX.pattern
@dataclass(frozen=True)
class EpisodeTarget:
season: int
episode: int
def normalize_title(title: str) -> str:
"""Lowercase, strip punctuation, collapse whitespace for fuzzy comparison."""
t = title.lower()
t = re.sub(r"[^\w\s]", " ", t, flags=re.UNICODE)
t = re.sub(r"\s+", " ", t).strip()
if t.startswith("the "):
t = t[4:]
return t
def _clean_title(rest: str) -> str:
title = rest.strip()
for prefix in ("- ", " ", "_ ", ". "):
if title.startswith(prefix):
title = title[len(prefix) :].strip()
if title.startswith("-") or title.startswith(""):
title = title[1:].strip()
if title.startswith("_"):
title = title[1:].strip()
return title
def parse_episode_stem(stem: str, pattern: str = DEFAULT_EPISODE_PATTERN) -> Optional[dict]:
"""
Parse common TV filename stems (S01E05 or 04x01 styles).
Returns dict with format, season, episode, title, padding hints — or None.
"""
del pattern # legacy param; auto-detect formats instead
m = PATTERN_SXXEXX.match(stem)
if m:
prefix_before_s, season_s, _e, ep_s, _dash, ep2_s, rest = m.groups()
try:
season = int(season_s)
old_first = int(ep_s)
except ValueError:
return None
span = 1
if ep2_s is not None:
try:
old_second = int(ep2_s)
except ValueError:
return None
span = old_second - old_first + 1
if span < 1:
span = 1
title = _clean_title(rest)
return {
"format": "sxxexx",
"prefix_before_s": prefix_before_s,
"season": season,
"season_pad": len(season_s),
"episode_pad": len(ep_s),
"old_first": old_first,
"span": span,
"title": title,
"suffix": rest,
}
m = PATTERN_NXNN.match(stem)
if m:
prefix, season_s, ep_s, sep, title_part = m.groups()
try:
season = int(season_s)
episode = int(ep_s)
except ValueError:
return None
title = _clean_title(title_part or "")
sep = sep or " "
if title and not sep.strip():
sep = " "
return {
"format": "nxnn",
"prefix": prefix,
"season": season,
"season_pad": len(season_s),
"episode_pad": len(ep_s),
"old_first": episode,
"span": 1,
"title": title,
"title_sep": sep,
}
return None
def rewrite_episode_stem(
stem: str,
target: EpisodeTarget,
padding: int = 2,
pattern: str = DEFAULT_EPISODE_PATTERN,
) -> str:
"""Replace season/episode block in stem, preserving layout and title."""
parsed = parse_episode_stem(stem, pattern)
if not parsed:
return stem
pad = max(1, padding)
new_season = target.season
new_ep = target.episode
title = parsed["title"]
span = parsed["span"]
if parsed["format"] == "nxnn":
s_pad = max(parsed["season_pad"], len(str(new_season)))
e_pad = max(parsed["episode_pad"], pad, len(str(new_ep)))
block = f"{new_season:0{s_pad}d}x{new_ep:0{e_pad}d}"
if title:
return f"{parsed['prefix']}{block}{parsed['title_sep']}{title}"
return f"{parsed['prefix']}{block}"
s_pad = max(parsed["season_pad"], len(str(new_season)))
e_pad = max(parsed["episode_pad"], pad, len(str(new_ep)))
e1 = str(new_ep).zfill(e_pad)
prefix = f"{parsed['prefix_before_s']}S{new_season:0{s_pad}d}E"
if span <= 1:
return f"{prefix}{e1}{parsed['suffix']}"
e2 = str(new_ep + span - 1).zfill(e_pad)
return f"{prefix}{e1}-E{e2}{parsed['suffix']}"
def rewrite_episode_number(
stem: str,
new_first_ep: int,
padding: int = 2,
pattern: str = DEFAULT_EPISODE_PATTERN,
) -> str:
"""Legacy helper: episode only, keep season from filename."""
parsed = parse_episode_stem(stem, pattern)
if not parsed:
return stem
return rewrite_episode_stem(
stem,
EpisodeTarget(season=parsed["season"], episode=new_first_ep),
padding=padding,
pattern=pattern,
)
def _similarity(a: str, b: str) -> float:
if not a or not b:
return 0.0
if a == b:
return 1.0
return SequenceMatcher(None, a, b).ratio()
def _coerce_target(value: EpisodeTarget | tuple[int, int] | int, parsed: dict) -> EpisodeTarget:
if isinstance(value, EpisodeTarget):
return value
if isinstance(value, tuple):
return EpisodeTarget(season=int(value[0]), episode=int(value[1]))
return EpisodeTarget(season=parsed["season"], episode=int(value))
def match_filenames_to_episodes(
filenames: list[str],
episodes: list[TvdbEpisode],
pattern: str = DEFAULT_EPISODE_PATTERN,
min_score: float = 0.65,
season_filter: int = 0,
) -> tuple[dict[str, EpisodeTarget], list[str], list[str]]:
"""
Match filenames to TheTVDB episodes by title.
season_filter: 0 = use all episodes; else only episodes from that season.
Returns mapping filename -> (season, episode), unmatched list, notes.
"""
if season_filter > 0:
episodes = [ep for ep in episodes if ep.season_number == season_filter]
file_entries: list[tuple[str, str, str, dict]] = []
for name in filenames:
base = Path(name).name
stem = base.rsplit(".", 1)[0] if "." in base and not base.startswith(".") else base
parsed = parse_episode_stem(stem, pattern)
if not parsed or not parsed["title"]:
continue
if season_filter > 0 and parsed["season"] != season_filter:
continue
norm = normalize_title(parsed["title"])
if norm:
file_entries.append((name, norm, parsed["title"], parsed))
ep_entries = [
(ep.season_number, ep.number, normalize_title(ep.name), ep.name)
for ep in episodes
]
pairs: list[tuple[float, str, int, int, str, str]] = []
for fname, fnorm, raw_title, _parsed in file_entries:
for season, ep_num, enorm, ep_name in ep_entries:
score = _similarity(fnorm, enorm)
pairs.append((score, fname, season, ep_num, raw_title, ep_name))
pairs.sort(key=lambda x: (-x[0], x[1], x[2], x[3]))
mapping: dict[str, EpisodeTarget] = {}
used_files: set[str] = set()
used_eps: set[tuple[int, int]] = set()
notes: list[str] = []
for score, fname, season, ep_num, raw_title, ep_name in pairs:
if score < min_score:
break
ep_key = (season, ep_num)
if fname in used_files or ep_key in used_eps:
continue
mapping[fname] = EpisodeTarget(season=season, episode=ep_num)
used_files.add(fname)
used_eps.add(ep_key)
pct = int(round(score * 100))
notes.append(
f"{fname}: S{season:02d}E{ep_num:02d} ← “{ep_name}” ({pct}% match, file title “{raw_title}”)"
)
unmatched_files = [
name for name in filenames
if name not in mapping
and parse_episode_stem(
(
Path(name).name.rsplit(".", 1)[0]
if "." in Path(name).name and not Path(name).name.startswith(".")
else Path(name).name
),
pattern,
)
]
return mapping, unmatched_files, notes