ffx/bin/ffx/track_descriptor.py

import logging

from .iso_language import IsoLanguage
from .track_type import TrackType
from .audio_layout import AudioLayout
from .track_disposition import TrackDisposition

from .helper import dictDiff, setDiff


class TrackDescriptor:

    CONTEXT_KEY = "context"

    ID_KEY = "id"
    INDEX_KEY = "index"
    SOURCE_INDEX_KEY = "source_index"
    SUB_INDEX_KEY = "sub_index"
    PATTERN_ID_KEY = "pattern_id"
    EXTERNAL_SOURCE_FILE_PATH_KEY = "external_source_file"

    DISPOSITION_SET_KEY = "disposition_set"
    TAGS_KEY = "tags"

    TRACK_TYPE_KEY = "track_type"
    CODEC_NAME_KEY = "codec_name"
    AUDIO_LAYOUT_KEY = "audio_layout"

    FFPROBE_INDEX_KEY = "index"
    FFPROBE_DISPOSITION_KEY = "disposition"
    FFPROBE_TAGS_KEY = "tags"
    FFPROBE_CODEC_TYPE_KEY = "codec_type"
    FFPROBE_CODEC_NAME_KEY = "codec_name"
    
    CODEC_PGS = 'hdmv_pgs_subtitle'

    def __init__(self, **kwargs):

        if TrackDescriptor.CONTEXT_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.CONTEXT_KEY]) is not dict:
                raise TypeError(
                    f"TrackDescriptor.__init__(): Argument {TrackDescriptor.CONTEXT_KEY} is required to be of type dict"
                )
            self.__context = kwargs[TrackDescriptor.CONTEXT_KEY]
            self.__logger = self.__context['logger']
        else:
            self.__context = {}
            self.__logger = logging.getLogger('FFX').addHandler(logging.NullHandler())

        if TrackDescriptor.ID_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.ID_KEY]) is not int:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.ID_KEY} is required to be of type int"
                )
            self.__trackId = kwargs[TrackDescriptor.ID_KEY]
        else:
            self.__trackId = -1

        if TrackDescriptor.PATTERN_ID_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.PATTERN_ID_KEY]) is not int:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.PATTERN_ID_KEY} is required to be of type int"
                )
            self.__patternId = kwargs[TrackDescriptor.PATTERN_ID_KEY]
        else:
            self.__patternId = -1

        if TrackDescriptor.EXTERNAL_SOURCE_FILE_PATH_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.EXTERNAL_SOURCE_FILE_PATH_KEY]) is not str:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.EXTERNAL_SOURCE_FILE_PATH_KEY} is required to be of type str"
                )
            self.__externalSourceFilePath = kwargs[TrackDescriptor.EXTERNAL_SOURCE_FILE_PATH_KEY]
        else:
            self.__externalSourceFilePath = ''

        if TrackDescriptor.INDEX_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.INDEX_KEY]) is not int:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.INDEX_KEY} is required to be of type int"
                )
            self.__index = kwargs[TrackDescriptor.INDEX_KEY]
        else:
            self.__index = -1

        if (
            TrackDescriptor.SOURCE_INDEX_KEY in kwargs.keys()
            and type(kwargs[TrackDescriptor.SOURCE_INDEX_KEY]) is int
        ):
            self.__sourceIndex = kwargs[TrackDescriptor.SOURCE_INDEX_KEY]
        else:
            self.__sourceIndex = self.__index

        if TrackDescriptor.SUB_INDEX_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.SUB_INDEX_KEY]) is not int:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.SUB_INDEX_KEY} is required to be of type int"
                )
            self.__subIndex = kwargs[TrackDescriptor.SUB_INDEX_KEY]
        else:
            self.__subIndex = -1

        if TrackDescriptor.TRACK_TYPE_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.TRACK_TYPE_KEY]) is not TrackType:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.TRACK_TYPE_KEY} is required to be of type TrackType"
                )
            self.__trackType = kwargs[TrackDescriptor.TRACK_TYPE_KEY]
        else:
            self.__trackType = TrackType.UNKNOWN

        if TrackDescriptor.CODEC_NAME_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.CODEC_NAME_KEY]) is not str:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.CODEC_NAME_KEY} is required to be of type str"
                )
            self.__codecName = kwargs[TrackDescriptor.CODEC_NAME_KEY]
        else:
            self.__codecName = ''

        if TrackDescriptor.TAGS_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.TAGS_KEY]) is not dict:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.TAGS_KEY} is required to be of type dict"
                )
            self.__trackTags = kwargs[TrackDescriptor.TAGS_KEY]
        else:
            self.__trackTags = {}

        if TrackDescriptor.DISPOSITION_SET_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.DISPOSITION_SET_KEY]) is not set:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.DISPOSITION_SET_KEY} is required to be of type set"
                )
            for d in kwargs[TrackDescriptor.DISPOSITION_SET_KEY]:
                if type(d) is not TrackDisposition:
                    raise TypeError(
                        f"TrackDesciptor.__init__(): All elements of argument set {TrackDescriptor.DISPOSITION_SET_KEY} is required to be of type TrackDisposition"
                    )
            self.__dispositionSet = kwargs[TrackDescriptor.DISPOSITION_SET_KEY]
        else:
            self.__dispositionSet = set()

        if TrackDescriptor.AUDIO_LAYOUT_KEY in kwargs.keys():
            if type(kwargs[TrackDescriptor.AUDIO_LAYOUT_KEY]) is not AudioLayout:
                raise TypeError(
                    f"TrackDesciptor.__init__(): Argument {TrackDescriptor.AUDIO_LAYOUT_KEY} is required to be of type AudioLayout"
                )
            self.__audioLayout = kwargs[TrackDescriptor.AUDIO_LAYOUT_KEY]
        else:
            self.__audioLayout = AudioLayout.LAYOUT_UNDEFINED

    @classmethod
    def fromFfprobe(cls, streamObj, subIndex: int = -1):
        """Processes ffprobe stream data as array with elements according to the following example
        {
        "index": 4,
        "codec_name": "hdmv_pgs_subtitle",
        "codec_long_name": "HDMV Presentation Graphic Stream subtitles",
        "codec_type": "subtitle",
        "codec_tag_string": "[0][0][0][0]",
        "codec_tag": "0x0000",
        "r_frame_rate": "0/0",
        "avg_frame_rate": "0/0",
        "time_base": "1/1000",
        "start_pts": 0,
        "start_time": "0.000000",
        "duration_ts": 1421035,
        "duration": "1421.035000",
        "disposition": {
                "default": 1,
                "dub": 0,
                "original": 0,
                "comment": 0,
                "lyrics": 0,
                "karaoke": 0,
                "forced": 0,
                "hearing_impaired": 0,
                "visual_impaired": 0,
                "clean_effects": 0,
                "attached_pic": 0,
                "timed_thumbnails": 0,
                "non_diegetic": 0,
                "captions": 0,
                "descriptions": 0,
                "metadata": 0,
                "dependent": 0,
                "still_image": 0
        },
        "tags": {
                "language": "ger",
                "title": "German Full"
        }
        }
        """

        trackType = (
            TrackType.fromLabel(streamObj["codec_type"])
            if "codec_type" in streamObj.keys()
            else TrackType.UNKNOWN
        )

        if trackType != TrackType.UNKNOWN:

            kwargs = {}

            kwargs[TrackDescriptor.INDEX_KEY] = (
                int(streamObj[TrackDescriptor.FFPROBE_INDEX_KEY])
                if TrackDescriptor.FFPROBE_INDEX_KEY in streamObj.keys()
                else -1
            )
            kwargs[TrackDescriptor.SOURCE_INDEX_KEY] = kwargs[TrackDescriptor.INDEX_KEY]
            kwargs[TrackDescriptor.SUB_INDEX_KEY] = subIndex

            kwargs[TrackDescriptor.TRACK_TYPE_KEY] = trackType
            kwargs[TrackDescriptor.CODEC_NAME_KEY] = str(streamObj[TrackDescriptor.FFPROBE_CODEC_NAME_KEY])

            kwargs[TrackDescriptor.DISPOSITION_SET_KEY] = (
                {
                    t
                    for d in (
                        k
                        for (k, v) in streamObj[
                            TrackDescriptor.FFPROBE_DISPOSITION_KEY
                        ].items()
                        if v
                    )
                    if (t := TrackDisposition.find(d)) is not None
                }
                if TrackDescriptor.FFPROBE_DISPOSITION_KEY in streamObj.keys()
                else set()
            )
            kwargs[TrackDescriptor.TAGS_KEY] = (
                streamObj[TrackDescriptor.FFPROBE_TAGS_KEY]
                if TrackDescriptor.FFPROBE_TAGS_KEY in streamObj.keys()
                else {}
            )
            kwargs[TrackDescriptor.AUDIO_LAYOUT_KEY] = (
                AudioLayout.identify(streamObj)
                if trackType == TrackType.AUDIO
                else AudioLayout.LAYOUT_UNDEFINED
            )

            return cls(**kwargs)
        else:
            return None

    def getId(self):
        return self.__trackId

    def getPatternId(self):
        return self.__patternId

    def getIndex(self):
        return self.__index

    def setIndex(self, index):
        self.__index = index

    def getSourceIndex(self):
        return self.__sourceIndex

    def setSourceIndex(self, sourceIndex: int):
        self.__sourceIndex = int(sourceIndex)

    def getSubIndex(self):
        return self.__subIndex

    def setSubIndex(self, subIndex):
        self.__subIndex = subIndex

    def getType(self):
        return self.__trackType
    
    def getCodec(self):
        return self.__codecName

    def getLanguage(self):
        if "language" in self.__trackTags.keys():
            return IsoLanguage.findThreeLetter(self.__trackTags["language"])
        else:
            return IsoLanguage.UNDEFINED

    def getTitle(self):
        if "title" in self.__trackTags.keys():
            return str(self.__trackTags["title"])
        else:
            return ""

    def getAudioLayout(self):
        return self.__audioLayout

    def getTags(self):
        return self.__trackTags

    def getDispositionSet(self):
        return self.__dispositionSet

    def getDispositionFlag(self, disposition: TrackDisposition) -> bool:
        return bool(disposition in self.__dispositionSet)

    def setDispositionFlag(self, disposition: TrackDisposition, state: bool):
        if state:
            self.__dispositionSet.add(disposition)
        else:
            self.__dispositionSet.discard(disposition)

    def compare(self, vsTrackDescriptor):

        compareResult = {}

        tagsDiffResult = dictDiff(vsTrackDescriptor.getTags(), self.getTags())

        if tagsDiffResult:
            compareResult[TrackDescriptor.TAGS_KEY] = tagsDiffResult

        vsDispositions = vsTrackDescriptor.getDispositionSet()
        dispositions = self.getDispositionSet()

        dispositionDiffResult = setDiff(vsDispositions, dispositions)

        if dispositionDiffResult:
            compareResult[TrackDescriptor.DISPOSITION_SET_KEY] = dispositionDiffResult

        return compareResult

    def setExternalSourceFilePath(self, filePath: str):
        self.__externalSourceFilePath = str(filePath)

    def getExternalSourceFilePath(self):
        return self.__externalSourceFilePath