From c818e69b0f1a8fc5cb34b97ca26cebe3374c3701 Mon Sep 17 00:00:00 2001 From: Jarian Date: Sun, 5 Jul 2026 08:09:04 +0000 Subject: [PATCH] feat: code quality improvements (#10, #12, #13) - Extract matching strategies into strategies.py module (#10) * DurationMatchStrategy: match files by runtime similarity * SequentialStrategy: fill remaining gaps in order * DiscMappingStrategy: explicit disc-to-episode mapping * ConstraintAwareStrategy: validate disc capacity + duration - Support multiple video formats: .mkv, .mp4, .avi, .mov, .wmv, .flv, .webm, .m4v (#12) - Make MediaInfo fallback ratio configurable via config.json duration_minutes_per_gb (#13) Closes #10, #12, #13 --- src/Matcher/__init__.py | 17 ++- src/Matcher/file_classifier.py | 57 ++++++---- src/Matcher/strategies.py | 200 +++++++++++++++++++++++++++++++++ src/config.py | 12 ++ 4 files changed, 266 insertions(+), 20 deletions(-) create mode 100644 src/Matcher/strategies.py diff --git a/src/Matcher/__init__.py b/src/Matcher/__init__.py index d9a4bd4..59f8674 100644 --- a/src/Matcher/__init__.py +++ b/src/Matcher/__init__.py @@ -1,5 +1,20 @@ """Episode Matcher module for classifying and renaming video files.""" from .file_classifier import FileClassifier from .episode_renamer import EpisodeRenamer +from .strategies import ( + DurationMatchStrategy, + SequentialStrategy, + DiscMappingStrategy, + ConstraintAwareStrategy, + match_episodes, +) -__all__ = ['FileClassifier', 'EpisodeRenamer'] +__all__ = [ + 'FileClassifier', + 'EpisodeRenamer', + 'DurationMatchStrategy', + 'SequentialStrategy', + 'DiscMappingStrategy', + 'ConstraintAwareStrategy', + 'match_episodes', +] diff --git a/src/Matcher/file_classifier.py b/src/Matcher/file_classifier.py index d454b3d..f0313e2 100644 --- a/src/Matcher/file_classifier.py +++ b/src/Matcher/file_classifier.py @@ -14,48 +14,67 @@ try: except ImportError: config_manager = None +# Supported video file extensions +DEFAULT_VIDEO_EXTENSIONS = {'.mkv', '.mp4', '.avi', '.mov', '.wmv', '.flv', '.webm', '.m4v'} + class FileClassifier: """Classifies video files as episodes or extras based on file size and video duration.""" - def __init__(self, folder_path: str): - """Initialize with folder path containing video files.""" + def __init__(self, folder_path: str, video_extensions: set = None): + """Initialize with folder path containing video files. + + Args: + folder_path: Path to folder containing video files. + video_extensions: Set of file extensions to scan (e.g., {'.mkv', '.mp4'}). + Defaults to DEFAULT_VIDEO_EXTENSIONS if not provided. + """ self.folder_path = Path(folder_path) + self.video_extensions = video_extensions or DEFAULT_VIDEO_EXTENSIONS self.video_files = self._get_video_files() self.file_info = self._analyze_files() def _get_video_files(self) -> List[Path]: - """Get all MKV files in the folder, excluding system files.""" + """Get all video files in the folder, excluding system files.""" if not self.folder_path.exists(): raise FileNotFoundError(f"Folder not found: {self.folder_path}") - - all_mkv_files = list(self.folder_path.glob("*.mkv")) + + video_files = [] + for ext in self.video_extensions: + video_files.extend(self.folder_path.glob(f"*{ext}")) # Filter out macOS resource fork files and other system files - video_files = [] - for file in all_mkv_files: + filtered = [] + for file in video_files: if file.name.startswith('._'): - continue # Skip macOS resource fork files + continue if file.name.startswith('.'): - continue # Skip any hidden files - video_files.append(file) + continue + filtered.append(file) - if not video_files: - print(f"Warning: No valid MKV files found in {self.folder_path}") + if not filtered: + ext_list = ', '.join(sorted(self.video_extensions)) + print(f"Warning: No video files found ({ext_list}) in {self.folder_path}") - return video_files + return filtered def _get_file_size(self, file_path: Path) -> int: """Get file size in bytes.""" return file_path.stat().st_size + def _get_fallback_duration_ratio(self) -> float: + """Get minutes-per-GB ratio for fallback duration estimation.""" + if config_manager: + return config_manager.get_duration_minutes_per_gb() + return 45.0 + def _get_video_duration(self, file_path: Path) -> float: """Get video duration in minutes using MediaInfo.""" + fallback_ratio = self._get_fallback_duration_ratio() if MediaInfo is None: - print("MediaInfo not available, using file size as proxy for duration") - # Rough estimate: 1GB ≈ 45 minutes for typical video + print(f"MediaInfo not available, using file size as proxy for duration ({fallback_ratio} min/GB)") size_gb = self._get_file_size(file_path) / (1024**3) - return size_gb * 45 + return size_gb * fallback_ratio try: media_info = MediaInfo.parse(str(file_path)) @@ -74,10 +93,10 @@ class FileClassifier: duration_minutes = float(duration_ms) / (1000 * 60) # Sanity check: if duration seems unreasonable, fall back to size estimation - if duration_minutes > 300: # More than 5 hours is likely wrong + if duration_minutes > 300: print(f"Warning: Suspicious duration ({duration_minutes:.1f} min) for {file_path.name}, using size estimation") size_gb = self._get_file_size(file_path) / (1024**3) - return size_gb * 45 + return size_gb * fallback_ratio return duration_minutes except Exception as e: @@ -85,7 +104,7 @@ class FileClassifier: # Fallback to size-based estimation size_gb = self._get_file_size(file_path) / (1024**3) - return size_gb * 45 + return size_gb * fallback_ratio def _analyze_files(self) -> List[Dict]: """Analyze all video files to get size and duration info.""" diff --git a/src/Matcher/strategies.py b/src/Matcher/strategies.py new file mode 100644 index 0000000..9cbfa68 --- /dev/null +++ b/src/Matcher/strategies.py @@ -0,0 +1,200 @@ +"""Matching strategies for episode-to-file mapping.""" +from pathlib import Path +from typing import List, Dict, Optional + + +class DurationMatchStrategy: + """Match files to TVDB episodes by duration similarity.""" + + def __init__(self, tolerance_minutes: float = 1.0): + self.tolerance = tolerance_minutes + + def match(self, files: List[Dict], tvdb_episodes: List[Dict], + allowed_episodes: Optional[List[int]] = None) -> List[Dict]: + matched = [] + used_files = set() + used_episodes = set() + + for file_info in files: + if id(file_info) in used_files: + continue + + file_duration = file_info['duration'] + best_match = None + best_score = float('inf') + + for tvdb_ep in tvdb_episodes: + if tvdb_ep['episode_number'] in used_episodes: + continue + if allowed_episodes and tvdb_ep['episode_number'] not in allowed_episodes: + continue + + tvdb_duration = tvdb_ep.get('runtime', 0) + if tvdb_duration == 0: + continue + + duration_diff = abs(file_duration - tvdb_duration) + if duration_diff < best_score: + best_score = duration_diff + best_match = tvdb_ep + + if best_match and best_score <= self.tolerance: + matched.append({ + 'file_info': file_info['file_info'], + 'episode_number': best_match['episode_number'], + 'tvdb_info': best_match, + 'duration_diff': best_score, + 'strategy': 'duration_match' + }) + used_files.add(id(file_info)) + used_episodes.add(best_match['episode_number']) + + return matched, used_files, used_episodes + + +class SequentialStrategy: + """Assign remaining files to episodes in sequential order.""" + + def match(self, remaining_files: List[Dict], tvdb_episodes: List[Dict], + used_episodes: set, allowed_episodes: List[int]) -> List[Dict]: + matched = [] + available = sorted( + [e for e in allowed_episodes if e not in used_episodes] + ) + + sorted_files = sorted(remaining_files, key=lambda x: x['filename']) + + for i, file_info in enumerate(sorted_files): + if i >= len(available): + break + + episode_number = available[i] + tvdb_match = next( + (ep for ep in tvdb_episodes if ep['episode_number'] == episode_number), + None + ) + + duration_diff = None + if tvdb_match and tvdb_match.get('runtime', 0) > 0: + duration_diff = abs(file_info['duration'] - tvdb_match['runtime']) + + matched.append({ + 'file_info': file_info['file_info'], + 'episode_number': episode_number, + 'tvdb_info': tvdb_match, + 'duration_diff': duration_diff, + 'strategy': 'sequential' + }) + + return matched + + +class DiscMappingStrategy: + """Match files to episodes using explicit disc-to-episode mapping.""" + + def __init__(self, disc_mapping: Dict[int, List[int]]): + self.disc_mapping = disc_mapping + + def match(self, episodes_info: List[Dict], tvdb_episodes: List[Dict]) -> List[Dict]: + matched = [] + disc_groups = {} + + for info in episodes_info: + disc_num = info['disc_number'] + if disc_num is not None: + disc_groups.setdefault(disc_num, []).append(info) + + for disc_num in sorted(disc_groups.keys()): + files = sorted(disc_groups[disc_num], key=lambda x: x['filename']) + expected = self.disc_mapping.get(disc_num, []) + + for i, file_info in enumerate(files): + if i >= len(expected): + break + + episode_number = expected[i] + tvdb_episode = next( + (ep for ep in tvdb_episodes if ep['episode_number'] == episode_number), + None + ) + + matched.append({ + 'file_info': file_info['file_info'], + 'episode_number': episode_number, + 'tvdb_info': tvdb_episode, + 'strategy': 'disc_mapping' + }) + + return matched + + +class ConstraintAwareStrategy: + """Match files respecting disc capacity and duration constraints.""" + + def __init__(self, max_duration_diff: float = 2.0, + min_file_duration_buffer: float = 2.0): + self.max_duration_diff = max_duration_diff + self.min_file_buffer = min_file_duration_buffer + + def validate(self, file_info: Dict, tvdb_episode: Dict, + disc_assignments: Dict[int, int], + disc_capacity: Dict[int, int]) -> bool: + file_duration = file_info['duration'] + tvdb_duration = tvdb_episode.get('runtime', 0) + file_disc = file_info['disc_number'] + + if tvdb_duration > 0: + diff = abs(file_duration - tvdb_duration) + if diff > self.max_duration_diff: + return False + if file_duration < (tvdb_duration - self.min_file_buffer): + return False + + if file_disc and disc_capacity: + max_eps = disc_capacity.get(file_disc, 0) + current = disc_assignments.get(file_disc, 0) + if current >= max_eps: + return False + + return True + + +def match_episodes(files: List[Dict], tvdb_episodes: List[Dict], + disc_mapping: Optional[Dict[int, List[int]]] = None, + disc_capacity: Optional[Dict[int, int]] = None + ) -> List[Dict]: + """ + Multi-strategy episode matcher. + + Pipeline: + 1. If disc_mapping provided, use DiscMappingStrategy + 2. Otherwise, use DurationMatchStrategy for close matches + 3. Fill gaps with SequentialStrategy + + Args: + files: Episode info dicts with file_info, disc_number, duration, filename + tvdb_episodes: TVDB episode data with runtime + disc_mapping: Optional disc-to-episode mapping + disc_capacity: Optional disc capacity constraints + + Returns: + List of match dicts sorted by episode number + """ + all_allowed = list(range(1, len(tvdb_episodes) + 1)) + + if disc_mapping: + return DiscMappingStrategy(disc_mapping).match(files, tvdb_episodes) + + # Phase 1: Duration matching + duration_strategy = DurationMatchStrategy() + matched, used_files, used_eps = duration_strategy.match( + files, tvdb_episodes, all_allowed + ) + + # Phase 2: Sequential fill for remaining + remaining = [f for f in files if id(f) not in used_files] + sequential = SequentialStrategy() + matched.extend(sequential.match(remaining, tvdb_episodes, used_eps, all_allowed)) + + matched.sort(key=lambda x: x['episode_number']) + return matched diff --git a/src/config.py b/src/config.py index 3714602..23ce3dc 100644 --- a/src/config.py +++ b/src/config.py @@ -23,6 +23,8 @@ class ConfigManager: default_config = { "tvdb_api_key": "", "default_episode_duration": 45, + "duration_minutes_per_gb": 45, + "video_extensions": [".mkv", ".mp4", ".avi", ".mov", ".wmv", ".flv", ".webm", ".m4v"], "classification_thresholds": { "size_threshold_ratio": 0.3, "duration_threshold_ratio": 0.4 @@ -65,6 +67,16 @@ class ConfigManager: """Get default episode duration in minutes.""" return self.config.get("default_episode_duration", 45) + def get_duration_minutes_per_gb(self) -> float: + """Get fallback duration estimation ratio (minutes per GB).""" + return float(self.config.get("duration_minutes_per_gb", 45)) + + def get_video_extensions(self) -> list: + """Get list of supported video file extensions.""" + return self.config.get("video_extensions", [ + ".mkv", ".mp4", ".avi", ".mov", ".wmv", ".flv", ".webm", ".m4v" + ]) + def get_classification_thresholds(self) -> Dict[str, float]: """Get classification threshold ratios.""" return self.config.get("classification_thresholds", {