Source code for multivoice.lib.stt_write_outputs

# src/multivoice/lib/stt_write_outputs.py

import os
import shutil
from multivoice.lib.stt_speaker_mapping import get_speaker_aware_transcript
from multivoice.lib.stt_timestamps import format_timestamp




def write_srt(transcript, file):
    """
    Write a transcript to a file in SRT format.

    Args:
        transcript (list): A list of dictionaries containing segment information.
        file (file object): The file object where the SRT content will be written.
    """
    for i, segment in enumerate(transcript, start=1):
        # write srt lines
        print(
            f"{i}\n"
            f"{format_timestamp(segment['start_time'], always_include_hours=True, decimal_marker=',')} --> "
            f"{format_timestamp(segment['end_time'], always_include_hours=True, decimal_marker=',')}\n"
            f"{segment['speaker']}: {segment['text'].strip().replace('-->', '->')}\n",
            file=file,
            flush=True,
        )





def cleanup(path: str):
    """path could either be relative or absolute."""
    # Check if the provided path is a file or directory and remove it accordingly
    if os.path.isfile(path) or os.path.islink(path):
        # Remove the file at the specified path
        os.remove(path)
    elif os.path.isdir(path):
        # Recursively delete the directory and all its contents
        shutil.rmtree(path)
    else:
        raise ValueError(f"Path {path} is not a file or dir.")





def write_outputs(ssm, args):
    """
    Write the speaker-aware transcript to both a text file and an SRT file.

    Args:
        ssm (object): An object containing the speech segment mapping.
        args (argparse.Namespace): Command-line arguments including output directory and audio file path.
    """
    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        # Create the output directory if it does not exist
        os.makedirs(output_dir)

    base_name = os.path.splitext(os.path.basename(args.audio))[0]
    txt_file_path = os.path.join(output_dir, f"{base_name}.txt")
    srt_file_path = os.path.join(output_dir, f"{base_name}.srt")

    with open(txt_file_path, "w", encoding="utf-8-sig") as f:
        get_speaker_aware_transcript(ssm, f)

    with open(srt_file_path, "w", encoding="utf-8-sig") as srt:
        write_srt(ssm, srt)