mm_avh_working_space/modules/mkvtoolnix.py

"""
    The MkvToolNix module provides a class for manipulating MKV files. It uses tools from the MKVToolNix package such as mkvextract, mkvmerge, and mkvinfo.

    * Example usage:
        mkvtoolnix = MkvToolNix(filename='example.mkv')
        mkv_info = mkvtoolnix.get_mkv_info()
        mkvtoolnix.mkv_extract_track(mkv_info)

    * Example output from get_mkv_info():
        {
            "container": {...},
            "global_tags": [...],
            "tracks": [...],
            "chapters": [...],
            "attachments": [...]
        }
"""

import sys
from subprocess import Popen, PIPE, CalledProcessError
from json import loads
from typing import Dict, List, Set
from os import path
from dataclasses import dataclass

from constants import (WORKING_SPACE,
                       WORKING_SPACE_OUTPUT,
                       WORKING_SPACE_TEMP,
                       MKV_EXTRACT_PATH,
                       MKV_MERGE_PATH, MKV_INFO_PATH,
                       MKV_PROPEDIT_PATH,
                       console)


@dataclass(slots=True)
class MkvToolNix:
    """
        A class for manipulating MKV files using the MKVToolNix package.

        Attributes:
            - filename (str): The name of the MKV file to be processed.
            - working_space (str): The directory where the MKV file is located.
            - working_space_output (str): The directory where the output files will be saved.
            - working_space_temp (str): The directory where temporary files will be saved during processing.
            - mkv_extract_path (str): The path to the mkvextract executable.
            - mkv_merge_path (str): The path to the mkvmerge executable.
            - mkv_info_path (str): The path to the mkvinfo executable.
            - mkv_propedit_path (str): The path to the mkvpropedit executable.

        Methods:
            - get_mkv_info(): Retrieves information about the MKV file using the mkvinfo tool.
            - mkv_extract_track(data: Dict[str, any]): Extracts the specified tracks from the MKV file using the mkvextract tool.
    """
    filename: str
    working_space: str = WORKING_SPACE
    working_space_output: str = WORKING_SPACE_OUTPUT
    working_space_temp: str = WORKING_SPACE_TEMP

    mkv_extract_path: str = MKV_EXTRACT_PATH
    mkv_merge_path: str = MKV_MERGE_PATH
    mkv_info_path: str = MKV_INFO_PATH
    mkv_propedit_path: str = MKV_PROPEDIT_PATH

    def _check_executables(self) -> None:
        """
            Checks if the MKVToolNix executables exist at the specified paths.
            If any executable is not found, the program will exit with an error message.
        """
        executables: List[str] = [self.mkv_extract_path,
                                  self.mkv_merge_path, self.mkv_info_path]
        for executable in executables:
            if not path.exists(executable):
                console.print(
                    f'Error: {executable} not found.', style='red_bold')
                sys.exit()

    def get_mkv_info(self) -> dict:
        """
            Retrieves information about the MKV file using the mkvinfo tool.
            The information is returned as a dictionary and also printed to the console.
            If an error occurs during the process, the program will exit with an error message.

            Returns:
            - dict: A dictionary containing information about the MKV file.
        """
        try:
            self._check_executables()
            command: List[str] = self._get_mkv_info_command()
            with Popen(command, stdout=PIPE, stderr=PIPE, universal_newlines=True) as process:
                output, error = process.communicate()

                if process.returncode == 0:
                    data: dict = loads(output)
                    tracks_data: List[dict] = self._parse_tracks_data(data)
                    self._print_mkv_info(tracks_data)
                    return data

                console.print(f'Error: {error}', style='red_bold')
        except (FileNotFoundError, CalledProcessError) as error:
            console.print(f'Error: {error}', style='red_bold')
            sys.exit()
        return {}

    def _get_mkv_info_command(self) -> List[str]:
        """
            Constructs the command to be used for retrieving information about the MKV file with mkvinfo.

            Returns:
                - List[str]: A list of strings representing the command to be executed.
        """
        return [
            self.mkv_merge_path,
            '--ui-language',
            'en',
            '--identify',
            '--identification-format',
            'json',
            path.join(self.working_space, self.filename)
        ]

    def _parse_tracks_data(self, data: dict) -> List[dict]:
        """
            Parses the track data from the dictionary returned by mkvinfo.

            Args:
                - data (dict): A dictionary containing information about the MKV file.

            Returns:
                - List[dict]: A list of dictionaries, each representing a track in the MKV file.
        """
        tracks_data: List[dict] = []

        for track in data['tracks']:
            track_data: dict = self._parse_track_data(track)
            tracks_data.append(track_data)

        return sorted(tracks_data, key=lambda x: x['id'])

    def _parse_track_data(self, track: dict) -> dict:
        """
            Parses the data for a single track from the dictionary returned by mkvinfo.
            Returns a dictionary with the track's ID, type, codec ID, language, IETF language, and properties.

            Args:
                - track (dict): A dictionary containing information about a single track in the MKV file.

            Returns:
                - dict: A dictionary containing information about a single track in the MKV file.
        """
        properties = track['properties']
        track_data: dict = {
            'id': track['id'],
            'type': track['type'],
            'codec_id': properties.get('codec_id', ''),
            'language': properties.get('language', ''),
            'language_ietf': properties.get('language_ietf', ''),
            'properties': self._get_track_properties(properties)
        }

        return track_data

    @staticmethod
    def _get_track_properties(properties: dict) -> str:
        """
            Retrieves the properties of a track from the dictionary returned by mkvinfo.

            Args:
                - properties (dict): A dictionary containing the properties of a track in the MKV file.

            Returns:
                - str: A string containing the properties of a track in the MKV file.
        """
        if 'display_dimensions' in properties:
            return properties['display_dimensions']
        if 'audio_sampling_frequency' in properties:
            return f"{properties['audio_sampling_frequency']} Hz"
        return 'None'

    def _print_mkv_info(self, tracks_data: List[dict]) -> None:
        """
            Prints the information about the MKV file to the console.
            The information includes the ID, type, codec ID, language, IETF language, and properties of each track.

            Args:
                - tracks_data (List[dict]): A list of dictionaries, each representing a track in the MKV file.
        """
        console.print('WYODRĘBNIANIE Z PLIKU:',
                      style='yellow_bold')
        console.print(self.filename, style='white_bold')
        console.print('ID  TYPE        CODEK                LANG  LANG_IETF  PROPERTIES',
                      style='yellow_bold')

        for track in tracks_data:
            console.print(
                f'[yellow_bold]{track["id"]:2}[/yellow_bold]  '
                f'[white_bold]{track["type"]:10}  '
                f'{track["codec_id"]:20} '
                f'{track["language"]:5} '
                f'{track["language_ietf"]:10} '
                f'{track["properties"]}'
            )
        console.print()

    def mkv_extract_track(self, data: Dict[str, any]) -> None:
        """
            Extracts the specified tracks from the MKV file using the mkvextract tool.
            The tracks to be extracted are specified by their IDs.
            The user is prompted to enter the IDs of the tracks to be extracted.
            If an error occurs during the process, the program will exit with an error message.

            Args:
                - data (Dict[str, any]): A dictionary containing information about the MKV file.
        """
        valid_track_range: range = range(len(data['tracks']))
        tracks_to_extract: Set[int] = set()

        while True:
            try:
                console.print('Podaj ID ścieżki do wyciągnięcia (naciśnij ENTER, aby zakończyć): ',
                              style='green_bold', end='')
                track_input: str = input().strip()
                if not track_input:
                    break
                track_id: int = int(track_input)

                if track_id in valid_track_range:
                    tracks_to_extract.add(track_id)
                else:
                    console.print(
                        'Nieprawidłowy ID ścieżki. Proszę podać poprawny numer ścieżki.\n', style='red_bold')
            except ValueError:
                console.print(
                    'Pominięto wyciąganie ścieżki.\n', style='red_bold')

        try:
            for track_id in tracks_to_extract:
                track: str = data['tracks'][track_id]
                codec_id: str = track['properties']['codec_id']
                format_extension: str = self._get_format_extension(codec_id)
                filename: str = f'{self.filename[:-4]}.{format_extension}'
                out_file: str = path.join(self.working_space_temp, filename)
                command: List[str] = self._get_extract_command(
                    track_id, out_file)

                with Popen(command) as process:
                    console.print(
                        f'\nEkstrakcja ścieżki {track_id} do pliku {filename}', style='yellow_bold')
                    process.wait()

        except (IndexError, KeyError):
            console.print(
                'Znaleziono nieprawidłowe ID ścieżki!', style='red_bold')
            self.mkv_extract_track(data)

        console.print(
            'Ekstrakcja zakończona pomyślnie.\n', style='green_bold')

    @staticmethod
    def _get_format_extension(codec_id: str) -> str:
        """
            Determines the file extension for a track based on its codec ID.

            Args:
                - codec_id (str): The codec ID of a track in the MKV file.

            Returns:
                - str: The file extension for the track.
        """
        format_dict: dict = {
            'A_AAC/MPEG2/*': 'aac',
            'A_AAC/MPEG4/*': 'aac',
            'A_AAC': 'aac',
            'A_AC3': 'ac3',
            'A_EAC3': 'ac3',
            'A_ALAC': 'caf',
            'A_DTS': 'dts',
            'A_FLAC': 'flac',
            'A_MPEG/L2': 'mp2',
            'A_MPEG/L3': 'mp3',
            'A_OPUS': 'opus',
            'A_PCM/INT/LIT': 'wav',
            'A_PCM/INT/BIG': 'wav',
            'A_REAL/*': 'rm',
            'A_TRUEHD': 'truehd',
            'A_MLP': 'mlp',
            'A_TTA1': 'tta',
            'A_VORBIS': 'ogg',
            'A_WAVPACK4': 'wv',
            'S_HDMV/PGS': 'sup',
            'S_HDMV/TEXTST': 'txt',
            'S_KATE': 'ogg',
            'S_TEXT/SSA': 'ssa',
            'S_TEXT/ASS': 'ass',
            'S_SSA': 'ssa',
            'S_ASS': 'ass',
            'S_TEXT/UTF8': 'srt',
            'S_TEXT/ASCII': 'srt',
            'S_VOBSUB': 'sub',
            'S_TEXT/USF': 'usf',
            'S_TEXT/WEBVTT': 'vtt',
            'V_MPEG1': 'mpeg',
            'V_MPEG2': 'mpeg',
            'V_MPEG4/ISO/AVC': 'h264',
            'V_MPEG4/ISO/HEVC': 'h265',
            'V_MS/VFW/FOURCC': 'avi',
            'V_REAL/*': 'rm',
            'V_THEORA': 'ogg',
            'V_VP8': 'ivf',
            'V_VP9': 'ivf'
        }

        return format_dict.get(codec_id, 'mkv')

    def _get_extract_command(self, track_id: int, out_file: str) -> List[str]:
        """
            Constructs the command to be used for extracting a track from the MKV file with mkvextract.
            Returns the command as a list of strings.

            Args:
                - track_id (int): The ID of the track to be extracted.
                - out_file (str): The path and filename of the output file.

            Returns:
                - List[str]: The command to be used for extracting the track.
        """
        return [
            self.mkv_extract_path,
            'tracks',
            path.join(self.working_space, self.filename),
            f'{track_id}:{out_file}'
        ]