mkv-episode-matcher 0.3.3__py3-none-any.whl → 1.0.0__py3-none-any.whl

mkv_episode_matcher/utils.py

@@ -2,15 +2,55 @@
 import os
 import re
 import shutil
+from pathlib import Path
 
 import requests
+import torch
 from loguru import logger
 from opensubtitlescom import OpenSubtitles
+from opensubtitlescom.exceptions import OpenSubtitlesException
+from rich.console import Console
+from rich.panel import Panel
 
-from mkv_episode_matcher.__main__ import CACHE_DIR, CONFIG_FILE
-from mkv_episode_matcher.config import get_config
+from mkv_episode_matcher.core.config_manager import get_config_manager
+from mkv_episode_matcher.subtitle_utils import find_existing_subtitle, sanitize_filename
 from mkv_episode_matcher.tmdb_client import fetch_season_details
 
+console = Console()
+
+
+def normalize_path(path_str):
+    """
+    Normalize a path string to handle cross-platform path issues.
+    Properly handles trailing slashes and backslashes in both Windows and Unix paths.
+    Also strips surrounding quotes that might be present in command line arguments.
+
+    Args:
+        path_str (str): The path string to normalize
+
+    Returns:
+        pathlib.Path: A normalized Path object
+    """
+    # Convert to string if it's a Path object
+    if isinstance(path_str, Path):
+        path_str = str(path_str)
+
+    # Strip surrounding quotes (both single and double)
+    path_str = path_str.strip().strip('"').strip("'")
+
+    # Remove trailing slashes or backslashes
+    path_str = path_str.rstrip("/").rstrip("\\")
+
+    # Handle Windows paths on non-Windows platforms
+    if os.name != "nt" and "\\" in path_str and ":" in path_str[:2]:
+        # This looks like a Windows path on a non-Windows system
+        # Extract the last component which should be the directory/file name
+        components = path_str.split("\\")
+        return Path(components[-1])
+
+    return Path(path_str)
+
+
 def get_valid_seasons(show_dir):
     """
     Get all season directories that contain MKV files.
@@ -22,39 +62,44 @@ def get_valid_seasons(show_dir):
         list: List of paths to valid season directories
     """
     # Get all season directories
-    season_paths = [
-        os.path.join(show_dir, d)
-        for d in os.listdir(show_dir)
-        if os.path.isdir(os.path.join(show_dir, d))
-    ]
+    show_path = normalize_path(show_dir)
+    season_paths = [str(show_path / d.name) for d in show_path.iterdir() if d.is_dir()]
 
     # Filter seasons to only include those with .mkv files
     valid_season_paths = []
     for season_path in season_paths:
-        mkv_files = [f for f in os.listdir(season_path) if f.endswith(".mkv")]
+        season_path_obj = Path(season_path)
+        mkv_files = [f for f in season_path_obj.iterdir() if f.name.endswith(".mkv")]
         if mkv_files:
             valid_season_paths.append(season_path)
 
     if not valid_season_paths:
-        logger.warning(f"No seasons with .mkv files found in show '{os.path.basename(show_dir)}'")
+        logger.warning(
+            f"No seasons with .mkv files found in show '{normalize_path(show_dir).name}'"
+        )
     else:
         logger.info(
-            f"Found {len(valid_season_paths)} seasons with .mkv files in '{os.path.basename(show_dir)}'"
+            f"Found {len(valid_season_paths)} seasons with .mkv files in '{normalize_path(show_dir).name}'"
         )
 
     return valid_season_paths
+
+
 def check_filename(filename):
     """
     Check if the filename is in the correct format (S01E02).
 
     Args:
-        filename (str): The filename to check.
+        filename (str or Path): The filename to check.
 
     Returns:
         bool: True if the filename matches the expected pattern.
     """
+    # Convert Path object to string if needed
+    if isinstance(filename, Path):
+        filename = str(filename)
     # Check if the filename matches the expected format
-    match = re.search(r'.*S\d+E\d+', filename)
+    match = re.search(r".*S\d+E\d+", filename)
     return bool(match)
 
 
@@ -70,16 +115,14 @@ def scramble_filename(original_file_path, file_number):
         None
     """
     logger.info(f"Scrambling {original_file_path}")
-    series_title = os.path.basename(
-        os.path.dirname(os.path.dirname(original_file_path))
-    )
-    original_file_name = os.path.basename(original_file_path)
-    extension = os.path.splitext(original_file_path)[-1]
+    series_title = normalize_path(original_file_path).parent.parent.name
+    original_file_name = Path(original_file_path).name
+    extension = Path(original_file_path).suffix
     new_file_name = f"{series_title} - {file_number:03d}{extension}"
-    new_file_path = os.path.join(os.path.dirname(original_file_path), new_file_name)
-    if not os.path.exists(new_file_path):
+    new_file_path = Path(original_file_path).parent / new_file_name
+    if not new_file_path.exists():
         logger.info(f"Renaming {original_file_name} -> {new_file_name}")
-        os.rename(original_file_path, new_file_path)
+        Path(original_file_path).rename(new_file_path)
 
 
 def rename_episode_file(original_file_path, new_filename):
@@ -87,32 +130,32 @@ def rename_episode_file(original_file_path, new_filename):
     Rename an episode file with a standardized naming convention.
 
     Args:
-        original_file_path (str): The original file path of the episode.
-        new_filename (str): The new filename including season/episode info.
+        original_file_path (str or Path): The original file path of the episode.
+        new_filename (str or Path): The new filename including season/episode info.
 
     Returns:
-        str: Path to the renamed file, or None if rename failed.
+        Path: Path to the renamed file, or None if rename failed.
     """
-    original_dir = os.path.dirname(original_file_path)
-    new_file_path = os.path.join(original_dir, new_filename)
-    
+    original_dir = Path(original_file_path).parent
+    new_file_path = original_dir / new_filename
+
     # Check if new filepath already exists
-    if os.path.exists(new_file_path):
+    if new_file_path.exists():
         logger.warning(f"File already exists: {new_filename}")
-        
+
         # Add numeric suffix if file exists
-        base, ext = os.path.splitext(new_filename)
+        base, ext = Path(new_filename).stem, Path(new_filename).suffix
         suffix = 2
         while True:
             new_filename = f"{base}_{suffix}{ext}"
-            new_file_path = os.path.join(original_dir, new_filename)
-            if not os.path.exists(new_file_path):
+            new_file_path = original_dir / new_filename
+            if not new_file_path.exists():
                 break
             suffix += 1
-    
+
     try:
-        os.rename(original_file_path, new_file_path)
-        logger.info(f"Renamed {os.path.basename(original_file_path)} -> {new_filename}")
+        Path(original_file_path).rename(new_file_path)
+        logger.info(f"Renamed {Path(original_file_path).name} -> {new_filename}")
         return new_file_path
     except OSError as e:
         logger.error(f"Failed to rename file: {e}")
@@ -120,28 +163,28 @@ def rename_episode_file(original_file_path, new_filename):
     except FileExistsError as e:
         logger.error(f"Failed to rename file: {e}")
         return None
-        
-def get_subtitles(show_id, seasons: set[int]):
+
+
+def get_subtitles(show_id, seasons: set[int], config=None, max_retries=3):
     """
     Retrieves and saves subtitles for a given TV show and seasons.
 
     Args:
         show_id (int): The ID of the TV show.
         seasons (Set[int]): A set of season numbers for which subtitles should be retrieved.
-
-    Returns:
-        None
+        config (Config object, optional): Preloaded configuration.
+        max_retries (int, optional): Number of times to retry subtitle download on OpenSubtitlesException. Defaults to 3.
     """
+    if config is None:
+        config = get_config_manager().load()
+    show_dir = config.show_dir
+    series_name = sanitize_filename(normalize_path(show_dir).name)
+    tmdb_api_key = config.tmdb_api_key
+    open_subtitles_api_key = config.open_subtitles_api_key
+    open_subtitles_user_agent = config.open_subtitles_user_agent
+    open_subtitles_username = config.open_subtitles_username
+    open_subtitles_password = config.open_subtitles_password
 
-    logger.info(f"Getting subtitles for show ID {show_id}")
-    config = get_config(CONFIG_FILE)
-    show_dir = config.get("show_dir")
-    series_name = os.path.basename(show_dir)
-    tmdb_api_key = config.get("tmdb_api_key")
-    open_subtitles_api_key = config.get("open_subtitles_api_key")
-    open_subtitles_user_agent = config.get("open_subtitles_user_agent")
-    open_subtitles_username = config.get("open_subtitles_username")
-    open_subtitles_password = config.get("open_subtitles_password")
     if not all([
         show_dir,
         tmdb_api_key,
@@ -151,92 +194,134 @@ def get_subtitles(show_id, seasons: set[int]):
         open_subtitles_password,
     ]):
         logger.error("Missing configuration settings. Please run the setup script.")
+        return
+
     try:
-        # Initialize the OpenSubtitles client
         subtitles = OpenSubtitles(open_subtitles_user_agent, open_subtitles_api_key)
-
-        # Log in (retrieve auth token)
         subtitles.login(open_subtitles_username, open_subtitles_password)
     except Exception as e:
         logger.error(f"Failed to log in to OpenSubtitles: {e}")
         return
+
     for season in seasons:
         episodes = fetch_season_details(show_id, season)
         logger.info(f"Found {episodes} episodes in Season {season}")
 
         for episode in range(1, episodes + 1):
             logger.info(f"Processing Season {season}, Episode {episode}...")
-            series_cache_dir = os.path.join(CACHE_DIR, "data", series_name)
+
+            series_cache_dir = config.cache_dir / "data" / series_name
             os.makedirs(series_cache_dir, exist_ok=True)
-            srt_filepath = os.path.join(
-                series_cache_dir,
-                f"{series_name} - S{season:02d}E{episode:02d}.srt",
-            )
-            if not os.path.exists(srt_filepath):
-                # get the episode info from TMDB
-                url = f"https://api.themoviedb.org/3/tv/{show_id}/season/{season}/episode/{episode}?api_key={tmdb_api_key}"
-                response = requests.get(url)
-                response.raise_for_status()
-                episode_data = response.json()
-                episode_data["name"]
-                episode_id = episode_data["id"]
-                # search for the subtitle
-                response = subtitles.search(tmdb_id=episode_id, languages="en")
-                if len(response.data) == 0:
-                    logger.warning(
-                        f"No subtitles found for {series_name} - S{season:02d}E{episode:02d}"
-                    )
-
-                for subtitle in response.data:
-                    subtitle_dict = subtitle.to_dict()
-                    # Remove special characters and convert to uppercase
-                    filename_clean = re.sub(
-                        r"\W+", " ", subtitle_dict["file_name"]
-                    ).upper()
-                    if f"E{episode:02d}" in filename_clean:
-                        logger.info(f"Original filename: {subtitle_dict['file_name']}")
-                        srt_file = subtitles.download_and_save(subtitle)
-                        series_name = series_name.replace(":", " -")
-                        shutil.move(srt_file, srt_filepath)
-                        logger.info(f"Subtitle saved to {srt_filepath}")
-                        break
-                    else:
-                        continue
-            else:
-                logger.info(
-                    f"Subtitle already exists for {series_name} - S{season:02d}E{episode:02d}"
-                )
-                continue
 
+            # Check for existing subtitle in any supported format
+            existing_subtitle = find_existing_subtitle(
+                series_cache_dir, series_name, season, episode
+            )
 
-def cleanup_ocr_files(show_dir):
-    """
-    Clean up OCR files generated during the episode matching process.
+            if existing_subtitle:
+                logger.info(f"Subtitle already exists: {Path(existing_subtitle).name}")
+                continue
 
-    Args:
-        show_dir (str): The directory containing the show files.
+            # Default to standard format for new downloads
+            srt_filepath = str(
+                series_cache_dir / f"{series_name} - S{season:02d}E{episode:02d}.srt"
+            )
 
-    Returns:
-        None
+            # get the episode info from TMDB
+            url = f"https://api.themoviedb.org/3/tv/{show_id}/season/{season}/episode/{episode}?api_key={tmdb_api_key}"
+            response = requests.get(url)
+            response.raise_for_status()
+            episode_data = response.json()
+            episode_id = episode_data["id"]
+
+            # search for the subtitle
+            response = subtitles.search(tmdb_id=episode_id, languages="en")
+            if len(response.data) == 0:
+                logger.warning(
+                    f"No subtitles found for {series_name} - S{season:02d}E{episode:02d}"
+                )
+                continue
 
-    This function cleans up the OCR files generated during the episode matching process.
-    It deletes the 'ocr' directory and all its contents in each season directory of the show.
-    """
-    for season_dir in os.listdir(show_dir):
-        season_dir_path = os.path.join(show_dir, season_dir)
-        ocr_dir_path = os.path.join(season_dir_path, "ocr")
-        if os.path.exists(ocr_dir_path):
-            logger.info(f"Cleaning up OCR files in {ocr_dir_path}")
-            shutil.rmtree(ocr_dir_path)
+            for subtitle in response.data:
+                subtitle_dict = subtitle.to_dict()
+                # Remove special characters and convert to uppercase
+                filename_clean = re.sub(
+                    r"\\W+", " ", subtitle_dict["file_name"]
+                ).upper()
+                if f"E{episode:02d}" in filename_clean:
+                    logger.info(f"Original filename: {subtitle_dict['file_name']}")
+                    retry_count = 0
+                    while retry_count < max_retries:
+                        try:
+                            srt_file = subtitles.download_and_save(subtitle)
+                            shutil.move(srt_file, srt_filepath)
+                            logger.info(f"Subtitle saved to {srt_filepath}")
+                            break
+                        except OpenSubtitlesException as e:
+                            retry_count += 1
+                            logger.error(
+                                f"OpenSubtitlesException (attempt {retry_count}): {e}"
+                            )
+                            console.print(
+                                f"[red]OpenSubtitlesException (attempt {retry_count}): {e}[/red]"
+                            )
+                            if retry_count >= max_retries:
+                                user_input = input(
+                                    "Would you like to continue matching? (y/n): "
+                                )
+                                if user_input.strip().lower() != "y":
+                                    logger.info(
+                                        "User chose to stop matching due to the error."
+                                    )
+                                    return
+                                else:
+                                    logger.info(
+                                        "User chose to continue matching despite the error."
+                                    )
+                                    break
+                        except Exception as e:
+                            logger.error(f"Failed to download and save subtitle: {e}")
+                            console.print(
+                                f"[red]Failed to download and save subtitle: {e}[/red]"
+                            )
+                            user_input = input(
+                                "Would you like to continue matching despite the error? (y/n): "
+                            )
+                            if user_input.strip().lower() != "y":
+                                logger.info(
+                                    "User chose to stop matching due to the error."
+                                )
+                                return
+                            else:
+                                logger.info(
+                                    "User chose to continue matching despite the error."
+                                )
+                                break
+                    else:
+                        continue
+                    break
 
 
 def clean_text(text):
-    # Remove brackets, parentheses, and their content
-    cleaned_text = re.sub(r"\[.*?\]|\(.*?\)|\{.*?\}", "", text)
+    # Remove brackets and curly braces with their content
+    cleaned_text = re.sub(r"\[.*?\]|\{.*?\}", "", text)
+    # Remove parentheses content EXCEPT for those containing exactly 4 digits (years)
+    # First, temporarily replace year patterns with a placeholder
+    year_pattern = r"\((\d{4})\)"
+    years = re.findall(year_pattern, cleaned_text)
+    cleaned_text = re.sub(year_pattern, "YEAR_PLACEHOLDER", cleaned_text)
+    # Remove all remaining parentheses content
+    cleaned_text = re.sub(r"\([^)]*\)", "", cleaned_text)
+    # Restore the years
+    for year in years:
+        cleaned_text = cleaned_text.replace("YEAR_PLACEHOLDER", f"({year})", 1)
+    # Normalize multiple spaces to single spaces
+    cleaned_text = re.sub(r"\s+", " ", cleaned_text)
     # Strip leading/trailing whitespace
     return cleaned_text.strip()
 
 
+@logger.catch
 def process_reference_srt_files(series_name):
     """
     Process reference SRT files for a given series.
@@ -248,24 +333,24 @@ def process_reference_srt_files(series_name):
         dict: A dictionary containing the reference files where the keys are the MKV filenames
               and the values are the corresponding SRT texts.
     """
-    from mkv_episode_matcher.__main__ import CACHE_DIR
-    import os
-    
+    config = get_config_manager().load()
+
     reference_files = {}
-    reference_dir = os.path.join(CACHE_DIR, "data", series_name)
-    
+    reference_dir = config.cache_dir / "data" / series_name
+
     for dirpath, _, filenames in os.walk(reference_dir):
         for filename in filenames:
             if filename.lower().endswith(".srt"):
-                srt_file = os.path.join(dirpath, filename)
+                srt_file = Path(dirpath) / filename
                 logger.info(f"Processing {srt_file}")
                 srt_text = extract_srt_text(srt_file)
                 season, episode = extract_season_episode(filename)
                 mkv_filename = f"{series_name} - S{season:02}E{episode:02}.mkv"
                 reference_files[mkv_filename] = srt_text
-                
+
     return reference_files
 
+
 def extract_srt_text(filepath):
     """
     Extracts text content from an SRT file.
@@ -279,39 +364,51 @@ def extract_srt_text(filepath):
     # Read the file content
     with open(filepath) as f:
         content = f.read()
-        
+
     # Split into subtitle blocks
-    blocks = content.strip().split('\n\n')
-    
+    blocks = content.strip().split("\n\n")
+
     text_lines = []
     for block in blocks:
-        lines = block.split('\n')
+        lines = block.split("\n")
         if len(lines) < 3:
             continue
-            
+
         # Skip index and timestamp, get all remaining lines as text
-        text = ' '.join(lines[2:])
+        text = " ".join(lines[2:])
         # Remove stage directions and tags
-        text = re.sub(r'\[.*?\]|\<.*?\>', '', text)
+        text = re.sub(r"\[.*?\]|\<.*?\>", "", text)
         if text:
             text_lines.append(text)
-            
+
     return text_lines
 
+
 def extract_season_episode(filename):
     """
-    Extract season and episode numbers from filename.
-    
+    Extract season and episode numbers from filename with support for multiple formats.
+
     Args:
         filename (str): Filename to parse
-        
+
     Returns:
         tuple: (season_number, episode_number)
     """
-    match = re.search(r'S(\d+)E(\d+)', filename)
-    if match:
-        return int(match.group(1)), int(match.group(2))
+    # List of patterns to try
+    patterns = [
+        r"S(\d+)E(\d+)",  # S01E01
+        r"(\d+)x(\d+)",  # 1x01 or 01x01
+        r"Season\s*(\d+).*?(\d+)",  # Season 1 - 01
+    ]
+
+    for pattern in patterns:
+        match = re.search(pattern, filename, re.IGNORECASE)
+        if match:
+            return int(match.group(1)), int(match.group(2))
+
     return None, None
+
+
 def process_srt_files(show_dir):
     """
     Process all SRT files in the given directory and its subdirectories.
@@ -326,11 +423,13 @@ def process_srt_files(show_dir):
     for dirpath, _, filenames in os.walk(show_dir):
         for filename in filenames:
             if filename.lower().endswith(".srt"):
-                srt_file = os.path.join(dirpath, filename)
+                srt_file = Path(dirpath) / filename
                 logger.info(f"Processing {srt_file}")
                 srt_text = extract_srt_text(srt_file)
                 srt_files[srt_file] = srt_text
     return srt_files
+
+
 def compare_and_rename_files(srt_files, reference_files, dry_run=False):
     """
     Compare the srt files with the reference files and rename the matching mkv files.
@@ -344,22 +443,21 @@ def compare_and_rename_files(srt_files, reference_files, dry_run=False):
         f"Comparing {len(srt_files)} srt files with {len(reference_files)} reference files"
     )
     for srt_text in srt_files.keys():
-        parent_dir = os.path.dirname(os.path.dirname(srt_text))
+        parent_dir = Path(srt_text).parent.parent
         for reference in reference_files.keys():
             _season, _episode = extract_season_episode(reference)
-            mkv_file = os.path.join(
-                parent_dir, os.path.basename(srt_text).replace(".srt", ".mkv")
-            )
+            mkv_file = str(parent_dir / Path(srt_text).name.replace(".srt", ".mkv"))
             matching_lines = compare_text(
                 reference_files[reference], srt_files[srt_text]
             )
             if matching_lines >= int(len(reference_files[reference]) * 0.1):
                 logger.info(f"Matching lines: {matching_lines}")
                 logger.info(f"Found matching file: {mkv_file} ->{reference}")
-                new_filename = os.path.join(parent_dir, reference)
+                new_filename = parent_dir / reference
                 if not dry_run:
-                    logger.info(f"Renaming {mkv_file} to {new_filename}")
-                    rename_episode_file(mkv_file, new_filename)
+                    logger.info(f"Renaming {mkv_file} to {str(new_filename)}")
+                    rename_episode_file(mkv_file, reference)
+
 
 def compare_text(text1, text2):
     """
@@ -378,4 +476,29 @@ def compare_text(text1, text2):
 
     # Compare the two lists of text lines
     matching_lines = set(flat_text1).intersection(flat_text2)
-    return len(matching_lines)
+    return len(matching_lines)
+
+
+def check_gpu_support():
+    logger.info("Checking GPU support...")
+    console.print("[bold]Checking GPU support...[/bold]")
+    if torch.cuda.is_available():
+        logger.info(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}")
+        console.print(
+            Panel.fit(
+                f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}",
+                title="GPU Support",
+                border_style="magenta",
+            )
+        )
+    else:
+        logger.warning(
+            "CUDA not available. Using CPU. Refer to https://pytorch.org/get-started/locally/ for GPU support."
+        )
+        console.print(
+            Panel.fit(
+                "CUDA not available. Using CPU. Refer to https://pytorch.org/get-started/locally/ for GPU support.",
+                title="GPU Support",
+                border_style="red",
+            )
+        )