MKV Subtitle Extraction Script

extract_subs.py

import subprocess
import json
import os
import sys
from concurrent.futures import ProcessPoolExecutor, as_completed

SUPPORTED_CODECS = {
    'S_TEXT/UTF8': 'srt',
    'S_TEXT/ASS': 'ass',
    'S_TEXT/SSA': 'ssa',
    'S_TEXT/USF': 'usf',
    'S_HDMV/PGS': 'sup'
}

def send_notification(title, body, percent=None, replace_id=None):
    """Sends or updates a GNOME notification."""
    cmd = ['notify-send', title, body, '-p'] # -p prints the ID
    
    if replace_id:
        cmd.extend(['-r', str(replace_id)]) # -r replaces existing notification
        
    if percent is not None:
        # GNOME hint for progress bar
        cmd.extend(['-h', f'int:value:{percent}'])
        # Make it transient (doesn't stay in history)
        cmd.extend(['-h', 'string:synchronous:extraction'])

    try:
        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
        return result.stdout.strip()
    except Exception:
        return None

def extract_subtitles(mkv_file):
    """Worker function to process a single MKV file."""
    try:
        # 1. Get track info
        result = subprocess.run(
            ['mkvmerge', '-J', mkv_file], 
            capture_output=True, text=True, check=True
        )
        data = json.loads(result.stdout)
        
        tracks = [t for t in data.get('tracks', []) if t['type'] == 'subtitles']
        
        if not tracks:
            return f"No subs: {os.path.basename(mkv_file)}"

        base_name = os.path.splitext(mkv_file)[0]
        extract_args = ['mkvextract', 'tracks', mkv_file]
        found_track = False

        for track in tracks:
            codec = track.get('properties', {}).get('codec_id')
            track_id = track['id']
            
            if codec in SUPPORTED_CODECS:
                ext = SUPPORTED_CODECS[codec]
                lang = track.get('properties', {}).get('language', 'und')
                out_file = f"{base_name}.{lang}.{ext}"
                
                if not os.path.exists(out_file):
                    extract_args.append(f"{track_id}:{out_file}")
                    found_track = True

        if found_track:
            subprocess.run(extract_args, check=True)
            return f"Extracted: {os.path.basename(mkv_file)}"
        
        return f"Skipped: {os.path.basename(mkv_file)}"

    except Exception as e:
        return f"Error: {os.path.basename(mkv_file)}"

def get_mkv_files(targets):
    """Scans arguments to build a unique list of MKV files."""
    mkv_list = []
    for target in targets:
        if os.path.isfile(target) and target.endswith(".mkv"):
            mkv_list.append(target)
        elif os.path.isdir(target):
            for f in os.listdir(target):
                if f.endswith(".mkv"):
                    mkv_list.append(os.path.join(target, f))
    return list(set(mkv_list))

if __name__ == "__main__":
    raw_args = sys.argv[1:] if len(sys.argv) > 1 else ["."]
    files_to_process = get_mkv_files(raw_args)
    total = len(files_to_process)

    if total == 0:
        send_notification("Subtitle Extractor", "No MKV files found.")
        sys.exit(0)

    # Initial Notification
    notif_id = send_notification(
        "Subtitle Extractor", 
        f"Starting extraction for {total} files...", 
        percent=0
    )

    with ProcessPoolExecutor() as executor:
        futures = {executor.submit(extract_subtitles, f): f for f in files_to_process}
        
        for i, future in enumerate(as_completed(futures), 1):
            result_msg = future.result()
            percent = int((i / total) * 100)
            
            # Update the SAME notification ID
            send_notification(
                "Subtitle Extractor", 
                f"{result_msg} ({i}/{total})", 
                percent=percent, 
                replace_id=notif_id
            )

    # Final success message (removes progress bar)
    send_notification(
        "Subtitle Extractor", 
        "All extractions complete!", 
        replace_id=notif_id
    )