diff --git a/main.py b/main.py index 5750e38..2516af7 100644 --- a/main.py +++ b/main.py @@ -1,17 +1,17 @@ +import sys import os +import traceback import re import logging +from concurrent.futures import ThreadPoolExecutor +from collections import defaultdict from aniworld.models import Anime, Episode from aniworld.common import get_season_episode_count, get_movie_episode_count from aniworld.search import search_anime from Loader import download -# Configure logging -logging.basicConfig( - filename="loader.log", - level=logging.DEBUG, - format="%(asctime)s - %(levelname)s - %(message)s" -) +# Configure basic logging to the console +logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') class MatchNotFoundError(Exception): """Custom exception raised when the pattern match is not found.""" @@ -20,18 +20,23 @@ class MatchNotFoundError(Exception): class Loader: def __init__(self, basePath: str): self.directory = basePath - logging.info(f"Initialized Loader with base path: {self.directory}") + logging.warning(f"Initialized Loader with base path: {self.directory}") def __find_mp4_files(self): - logging.info("Scanning for .mp4 files") - for root, dirs, files in os.walk(self.directory): - mp4_files = [file for file in files if file.endswith('.mp4')] + logging.warning("Scanning for .mp4 files") + + for root_folder_name in os.listdir(self.directory): + + folder_data = defaultdict(list) # Dictionary to store MP4 files per folder + folder = os.path.join(self.directory, root_folder_name) + + # First pass: Scan all folders and collect MP4 file data + for root, dirs, files in os.walk(folder): + mp4_files = [file for file in files if file.endswith('.mp4')] + if mp4_files: + folder_data[root_folder_name].extend(mp4_files) + yield root_folder_name, folder_data[root_folder_name] - if mp4_files: - relative_path = os.path.relpath(root, self.directory) - root_folder_name = relative_path.split(os.sep)[0] - logging.debug(f"Found {len(mp4_files)} .mp4 files in {root_folder_name}") - yield root_folder_name, mp4_files for dir in self.__find_empty_folders(): logging.debug(f"Found no .mp4 files in {dir}") yield dir, [] @@ -62,14 +67,17 @@ class Loader: logging.debug(f"Key found for folder '{folder_name}': {key}") return key else: - key = search_anime(self.__remove_year(folder_name)) - with open(key_file, 'w') as file: - file.write(key) - logging.info(f"Generated new key for folder '{folder_name}': {key}") - return key + key = search_anime(self.__remove_year(folder_name), True) + if (len(key) >= 1): + key = key[0]['link'] + with open(key_file, 'w') as file: + file.write(key) + logging.warning(f"Generated new key for folder '{folder_name}': {key}") + return key + raise Exception() def __GetEpisodeAndSeason(self, filename: str): - pattern = r'S(\d{2})E(\d{2})' + pattern = r'S(\d+)E(\d+)' match = re.search(pattern, filename) if match: season = match.group(1) @@ -104,36 +112,42 @@ class Loader: if missing_episodes: yield season, missing_episodes - def LoadMissing(self): - logging.info("Starting process to load missing episodes") + logging.warning("Starting process to load missing episodes") result = self.__find_mp4_files() - for folder, mp4_files in result: + def download_episode(folder, season, episode, key): + """Helper function to download an individual episode.""" try: - key = self.__check_and_generate_key(folder) - missings = self.__GetMissingEpisodesAndSeason(key, mp4_files) - - for season, missing_episodes in missings: - folder_path = os.path.join(self.directory, folder, f"Season {season}") - - for episode in missing_episodes: - anime = Anime( - episode_list=[ - Episode(slug=key, season=season, episode=episode) - ], - language="German Dub", - output_directory=folder_path - ) - logging.info(f"Downloading episode {episode} of season {season} for anime {key}") - download(anime) + folder_path = os.path.join(self.directory, folder, f"Season {season}") + anime = Anime( + episode_list=[Episode(slug=key, season=season, episode=episode)], + language="German Dub", + output_directory=folder_path + ) + logging.warning(f"Downloading episode {episode} of season {season} for anime {key}") + download(anime) except Exception as e: - logging.error(f"Error processing folder '{folder}': {e}") - continue + logging.error(f"Error downloading episode {episode} of season {season} for anime {key}: {e}") + + # Using ThreadPoolExecutor to run downloads in parallel + with ThreadPoolExecutor(max_workers=5) as executor: # Adjust number of workers as needed + for folder, mp4_files in result: + try: + key = self.__check_and_generate_key(folder) + missings = self.__GetMissingEpisodesAndSeason(key, mp4_files) + + for season, missing_episodes in missings: + for episode in missing_episodes: + executor.submit(download_episode, folder, season, episode, key) + except Exception as e: + logging.error(f"Error processing folder '{folder}': {e}") + traceback.print_exc() + continue # Read the base directory from an environment variable -#directory_to_search = os.getenv("ANIME_DIRECTORY", "\\\\sshfs.r\\ubuntu@192.168.178.43\\media\\serien\\Serien") -directory_to_search = os.getenv("ANIME_DIRECTORY", "D:\sss") +directory_to_search = os.getenv("ANIME_DIRECTORY", "\\\\sshfs.r\\ubuntu@192.168.178.43\\media\\serien\\Serien") +#directory_to_search = os.getenv("ANIME_DIRECTORY", "D:\sss") loader = Loader(directory_to_search) loader.LoadMissing() \ No newline at end of file