From cf3e4acc43939c1402bb5b91e0196839391b87c7 Mon Sep 17 00:00:00 2001 From: Karl Date: Tue, 23 Jul 2024 15:09:01 +0000 Subject: [PATCH] docstrings --- bazarr-ai-sub-generator/main.py | 77 ++++++++++++++++++++++++++++++--- 1 file changed, 71 insertions(+), 6 deletions(-) diff --git a/bazarr-ai-sub-generator/main.py b/bazarr-ai-sub-generator/main.py index 5021f5a..7ad2eb0 100644 --- a/bazarr-ai-sub-generator/main.py +++ b/bazarr-ai-sub-generator/main.py @@ -2,6 +2,7 @@ import os import warnings import tempfile import time +from typing import List, Dict, Any from utils.files import filename, write_srt from utils.ffmpeg import get_audio, add_subtitles_to_mp4 from utils.bazarr import get_wanted_episodes, get_episode_details, sync_series @@ -11,7 +12,18 @@ from utils.whisper import WhisperAI from utils.decorator import measure_time -def process_audio_and_subtitles(file_path, model_args, args, backend): +def process_audio_and_subtitles(file_path: str, model_args: Dict[str, Any], args: Dict[str, Any], backend: str) -> None: + """Processes audio extraction and subtitle generation for a given file. + + Args: + file_path (str): Path to the video file. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + args (Dict[str, Any]): Additional arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + None + """ try: audios = get_audio([file_path], 0, None) subtitles = get_subtitles(audios, tempfile.gettempdir(), model_args, args, backend) @@ -21,7 +33,18 @@ def process_audio_and_subtitles(file_path, model_args, args, backend): print(f"Skipping file {file_path} due to - {ex}") -def folder_flow(folder, model_args, args, backend): +def folder_flow(folder: str, model_args: Dict[str, Any], args: Dict[str, Any], backend: str) -> None: + """Processes all files within a specified folder. + + Args: + folder (str): Path to the folder containing video files. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + args (Dict[str, Any]): Additional arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + None + """ print(f"Processing folder {folder}") files = os.listdir(folder) for file in files: @@ -30,12 +53,34 @@ def folder_flow(folder, model_args, args, backend): process_audio_and_subtitles(path, model_args, args, backend) -def file_flow(file_path, model_args, args, backend): +def file_flow(file_path: str, model_args: Dict[str, Any], args: Dict[str, Any], backend: str) -> None: + """Processes a single specified file. + + Args: + file_path (str): Path to the video file. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + args (Dict[str, Any]): Additional arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + None + """ print(f"Processing file {file_path}") process_audio_and_subtitles(file_path, model_args, args, backend) -def bazzar_flow(show, model_args, args, backend): +def bazzar_flow(show: str, model_args: Dict[str, Any], args: Dict[str, Any], backend: str) -> None: + """Processes episodes needing subtitles from Bazarr API. + + Args: + show (str): The show name. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + args (Dict[str, Any]): Additional arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + None + """ list_of_episodes_needing_subtitles = get_wanted_episodes(show) print(f"Found {list_of_episodes_needing_subtitles['total']} episodes needing subtitles.") for episode in list_of_episodes_needing_subtitles["data"]: @@ -47,7 +92,19 @@ def bazzar_flow(show, model_args, args, backend): @measure_time -def get_subtitles(audio_paths: list, output_dir: str, model_args: dict, transcribe_args: dict, backend: str): +def get_subtitles(audio_paths: List[str], output_dir: str, model_args: Dict[str, Any], transcribe_args: Dict[str, Any], backend: str) -> Dict[str, str]: + """Generates subtitles for given audio files using the specified model. + + Args: + audio_paths (List[str]): List of paths to the audio files. + output_dir (str): Directory to save the generated subtitle files. + model_args (Dict[str, Any]): Model arguments for subtitle generation. + transcribe_args (Dict[str, Any]): Transcription arguments for subtitle generation. + backend (str): Backend to use ('whisper' or 'faster_whisper'). + + Returns: + Dict[str, str]: A dictionary mapping audio file paths to generated subtitle file paths. + """ if backend == 'whisper': model = WhisperAI(model_args, transcribe_args) else: @@ -68,7 +125,15 @@ def get_subtitles(audio_paths: list, output_dir: str, model_args: dict, transcri return subtitles_path -def process(args: dict): +def process(args: Dict[str, Any]) -> None: + """Main entry point to determine which processing flow to use. + + Args: + args (Dict[str, Any]): Dictionary of arguments including model, language, show, file, folder, and backend. + + Returns: + None + """ model_name: str = args.pop("model") language: str = args.pop("language") show: str = args.pop("show")