bootylift.py

'''
bootylift: it fixes ASS
-iterates recursively through given directory
-checks every mkv file for SubStationAlpha subtitles
-extracts all ass/ssa subtitles
-converts them to srt
-creates new mkv file without the ass subs replaced by the new srt subs
-cleans up after itself
'''

import os
import re
import sys
import subprocess
import pymediainfo

def cleanup(file, sub_ids):
    subprocess.call(["rm", file])
    subprocess.call(["mv", file[:-4]+'-temp.mkv', file])
    for elem in sub_ids:
        subprocess.call(["rm", file[:-4]+'-'+str(elem[0])+'.srt'])
    print("Finished handling", file, "–", str(len(sub_ids)), "subtitles tracks converted from ASS/SSA to SRT.")

def remux(file, sub_ids):
    command = ["mkvmerge", "-o", file[:-4]+'-temp.mkv', '--no-subtitles', file]
    for elem in sub_ids:
        elem_id = elem[0]
        lang = elem[1]
        if lang == None:
            lang = 'und'
        srt_filename = file[:-4]+'-'+str(elem_id)+'.srt'
        command.append("--language")
        command.append('0:'+lang)
        command.append(srt_filename)
    subprocess.call(command)

def convert_ass(file, sub_ids, root):
    for elem in sub_ids:
        ass_filename = file[:-4]+'-'+str(elem[0])+'.ass'
        subprocess.call(["asstosrt", ass_filename, "--output-dir", root])
        subprocess.call(["rm", file[:-4]+'-'+str(elem[0])+'.ass'])

def extract_ass(file, sub_ids):
    for elem in sub_ids:
        ass_filename = file[:-4]+'-'+str(elem[0])+'.ass'
        subprocess.call(["mkvextract", "tracks", file, str(elem[0])+':'+ass_filename])

def fix_mkv(root, file, sub_ids):
    extract_ass(file, sub_ids)
    convert_ass(file, sub_ids, root)
    remux(file, sub_ids)
    cleanup(file, sub_ids)

def inspect_mkv(root, file):
    '''
    Runs an mkvmerge analysis on the mkv and finds all ASS/SSA subs, collecting their track IDs and language tags
    '''
    proc = subprocess.Popen(["mkvmerge", "-i", file], stdout=subprocess.PIPE)
    output = proc.stdout.read().decode()
    ass_subs = re.findall('Track ID \d: subtitles \(SubStationAlpha\)', output)
    sub_ids = list()
    for elem in ass_subs:
        track_id = int(re.search('Track ID (\d)', elem).group(1))
        lang = pymediainfo.MediaInfo.parse(file).tracks[track_id].language
        sub_ids.append((track_id, lang))
    if len(sub_ids) > 0:
        fix_mkv(root, file, sub_ids)


def main():
    dir = sys.argv[1]
    for root, dirs, files in os.walk(dir):
        for file in files:
            if file.endswith(".mkv"):
                filepath = os.path.join(root, file)
                inspect_mkv(root, filepath)

if __name__ == "__main__":
    main()