Скрипт позволяющий смёржить в mkv видео аудиодорожки и субтитры, выбрав при этом приоритетные (для VLC/PLEX)

mkv_merger.py

#!/usr/bin/env python
import os
import re
import subprocess
from os import listdir
from os.path import isfile, join
from typing import Optional

"""
Скрипт позволяет объединить видео mkv, аудиодорожки и субтитры в один mkv контейнер. Берёт все аудиодорожки и субтитры
которые найдёт во вложенных директориях, имена которых совпадут с названием mkv файла
Умеет автоматически проставлять названия субтитрам и выбирать язык
Умеет проставлять FORCED_DISPLAY для дорожек

1. Поместить скрипт в директорию с mkv в корневой директории и с аудиодорожками/субтитрами во вложенных директориях
2. Запустить скрипт
3. Выбрать форсированные дорожки, если требуется
"""

MKV = "mkv"
M4A = "m4a"
MKA = "mka"
ASS = "ass"
SRT = "srt"

CUR_DIR = os.getcwd()

mkv_files = {}
all_audio_tracks = set()
all_subtitles = set()


class MKVFile:
    MKV_MERGE = r"C:\Program Files\MKVToolNix\mkvmerge.exe"
    TRACK_ID_RE = r"Track ID (\d*)"

    def __init__(self, name: str):
        self.name: str = name
        self.video_track: str = ""
        self.audio_tracks: list[str] = []
        self.subtitles: list[str] = []

        self.forced_audio_index = None
        self.forced_subtitle_index = None

    @property
    def output_filename(self):
        return os.path.join(CUR_DIR, f"{self.name}_merged.{MKV}")

    @staticmethod
    def _get_lang(full_name: str):
        """
        Кривая попытка получить язык по полному пути файла
        """
        full_name_lower = full_name.lower()

        language = "und"
        if "ru" in full_name_lower:
            language = "ru"
        elif "en" in full_name_lower:
            language = "en"
        return language

    def _get_mkv_tracks(self) -> list[str]:
        """
        Получение информации о mkv видео. Вытаскивает список всех дорожек
        """

        cmd = [f'"{self.MKV_MERGE}"', "-i", f'"{self.video_track}"']
        cmd = " ".join(cmd)

        tracks = []
        file_id = 0
        with subprocess.Popen(cmd, stdout=subprocess.PIPE, universal_newlines=True) as process:
            for line in process.stdout:
                match = re.match(self.TRACK_ID_RE, line)
                if match:
                    track_id = match.group(1)
                    track = f"{file_id}:{track_id}"
                    tracks.append(track)
        return tracks

    def _get_tracks_cmd(self, tracks, forced_track, forced_var) -> list[str]:
        """
        Формирование дорожек аудио или субтитров. Проставление имён и языка. Выбор forced, если есть
        """
        tracks_cmd = []
        if tracks:
            track_id = 0
            for i, file in enumerate(tracks):
                full_name = "_".join(file.replace(CUR_DIR, "").strip("\\").split('\\')[:-1])
                language = self._get_lang(full_name)
                tracks_cmd.append(f'--language {track_id}:{language}')
                tracks_cmd.append(f'--track-name {track_id}:"{full_name}"')
                if full_name == forced_track:
                    tracks_cmd.append(f'--forced-display-flag {track_id}:yes')
                    setattr(self, forced_var, i)
                tracks_cmd.append(f'"{file}"')
        return tracks_cmd

    def _get_tracks_order(self) -> str:
        """
        Выстраивание дорожек в нужном порядке.
        Подробнее читай коменты
        """

        mkv_tracks = self._get_mkv_tracks()
        audio_tracks = [f"{i + 1}:0" for i, _ in enumerate(self.audio_tracks)]
        subtitles = [f"{i + 1 + len(self.audio_tracks)}:0" for i, _ in enumerate(self.subtitles)]

        forced_audio_track_id = None
        forced_subtitle_id = None

        if self.forced_audio_index is not None:
            forced_audio_track_id = audio_tracks[self.forced_audio_index]
            audio_tracks = audio_tracks[:self.forced_audio_index] + audio_tracks[self.forced_audio_index + 1:]
        if self.forced_subtitle_index is not None:
            forced_subtitle_id = subtitles[self.forced_subtitle_index]
            subtitles = subtitles[:self.forced_subtitle_index] + subtitles[self.forced_subtitle_index + 1:]

        # Сначала всегда идёт видео
        tracks_order = [mkv_tracks[0]]
        # forced аудиодорожка идёт следом
        if forced_audio_track_id:
            tracks_order.append(forced_audio_track_id)
        # forced субтитры идёт следом
        if forced_subtitle_id:
            tracks_order.append(forced_subtitle_id)
        # Остальные аудиодорожки
        if audio_tracks:
            tracks_order += audio_tracks
        # Остальные субтитры
        if subtitles:
            tracks_order += subtitles
        # Оригинальные дорожки видео
        tracks_order += mkv_tracks[1:]

        return f"--track-order {','.join(tracks_order)}"

    def get_merge_command(self, forced_audio=None, forced_subtitle=None) -> Optional[str]:
        if not self.audio_tracks and not self.subtitles:
            return None

        # Вызов команды, указание выходного файла, указание видеодорожки
        cmd = [f'"{self.MKV_MERGE}"', "-o", f'"{self.output_filename}"', f'"{self.video_track}"']

        # Указание аудиодорожек
        if audio_tracks := self._get_tracks_cmd(self.audio_tracks, forced_audio, "forced_audio_index"):
            cmd += audio_tracks
        # Указание дорожек субтитров
        if subtitles := self._get_tracks_cmd(self.subtitles, forced_subtitle, "forced_subtitle_index"):
            cmd += subtitles

        # Упорядочивание
        cmd.append(self._get_tracks_order())

        return " ".join(cmd)


def get_video_audio_subtitles():
    # Сканируем CUR_DIR и запоминаем названия всех mkv файлов
    cur_dir_files = [f for f in listdir(CUR_DIR) if isfile(join(CUR_DIR, f))]
    for file in cur_dir_files:
        try:
            filename, ext = file.rsplit('.', 1)
        except ValueError:
            continue

        if ext == MKV:
            mkv_files[filename] = MKVFile(filename)
            mkv_files[filename].video_track = os.path.join(CUR_DIR, file)

    # Обход всех вложенных директорий для поиска аудиодорожек и субтитров
    for dirname, _, files in os.walk(CUR_DIR):
        if not files:
            continue
        for file in files:
            try:
                filename, ext = file.rsplit('.', 1)
            except ValueError:
                continue

            if filename not in mkv_files:
                find = False
                for mkv_file in mkv_files.keys():
                    if mkv_file in filename:
                        filename = mkv_file
                        find = True
                if not find:
                    continue

            if ext in [M4A, MKA]:
                full_file_path = os.path.join(dirname, file)
                full_name = "_".join(full_file_path.replace(CUR_DIR, "").strip("\\").split('\\')[:-1])

                mkv_files[filename].audio_tracks.append(full_file_path)
                all_audio_tracks.add(full_name)
            elif ext in [ASS, SRT]:
                full_file_path = os.path.join(dirname, file)
                full_name = "_".join(full_file_path.replace(CUR_DIR, "").strip("\\").split('\\')[:-1])

                mkv_files[filename].subtitles.append(full_file_path)
                all_subtitles.add(full_name)
    return mkv_files


def request_forced_info(entities, what_to_request) -> Optional[str]:
    """
    Запрос информации о выбранных дорожках (неважно каких)
    """
    entities = {str(i + 1): name for i, name in enumerate(sorted(entities))}
    if entities:
        entities_str = "\n".join([f"{key}. {value}" for key, value in entities.items()])

        print(f"Выберите FORCED дорожку {what_to_request}:\n{entities_str}\n\nЕсли не нужно - нажмите Enter")
        entity_id = input()

        try:
            return entities[entity_id]
        except KeyError:
            return None


def main():
    mkv_files = get_video_audio_subtitles()
    forced_audio = request_forced_info(all_audio_tracks, "аудио")
    forced_subtitle = request_forced_info(all_subtitles, "субтитров")
    commands = [item.get_merge_command(forced_audio, forced_subtitle) for item in mkv_files.values()]

    for command in commands:
        if command:
            subprocess.run(command)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass