From 057e4efb566b126ed647f29d4783911630e6d9b2 Mon Sep 17 00:00:00 2001 From: rlaphoenix Date: Mon, 25 Mar 2024 05:38:32 +0000 Subject: [PATCH] feat: Add support for MKV Attachments via Attachment class You add these new Attachment objects to the Tracks object just like you would with Video, Audio, and Subtitle objects. --- devine/commands/dl.py | 2 +- devine/core/tracks/attachment.py | 70 ++++++++++++++++++++++++++++++++ devine/core/tracks/tracks.py | 55 +++++++++++++++++++++---- 3 files changed, 117 insertions(+), 10 deletions(-) create mode 100644 devine/core/tracks/attachment.py diff --git a/devine/commands/dl.py b/devine/commands/dl.py index f853d30..1633cc0 100644 --- a/devine/commands/dl.py +++ b/devine/commands/dl.py @@ -632,7 +632,7 @@ class dl: task_id = progress.add_task(f"{task_description}...", total=None, start=False) - task_tracks = Tracks(title.tracks) + title.tracks.chapters + task_tracks = Tracks(title.tracks) + title.tracks.chapters + title.tracks.attachments if video_track: task_tracks.videos = [video_track] diff --git a/devine/core/tracks/attachment.py b/devine/core/tracks/attachment.py new file mode 100644 index 0000000..c359629 --- /dev/null +++ b/devine/core/tracks/attachment.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import mimetypes +from pathlib import Path +from typing import Optional, Union +from zlib import crc32 + + +class Attachment: + def __init__( + self, + path: Union[Path, str], + name: Optional[str] = None, + mime_type: Optional[str] = None, + description: Optional[str] = None + ): + """ + Create a new Attachment. + + If name is not provided it will use the file name (without extension). + If mime_type is not provided, it will try to guess it. + """ + if not isinstance(path, (str, Path)): + raise ValueError("The attachment path must be provided.") + if not isinstance(name, (str, type(None))): + raise ValueError("The attachment name must be provided.") + + path = Path(path) + if not path.exists(): + raise ValueError("The attachment file does not exist.") + + name = (name or path.stem).strip() + mime_type = (mime_type or "").strip() or None + description = (description or "").strip() or None + + if not mime_type: + mime_type = { + ".ttf": "application/x-truetype-font", + ".otf": "application/vnd.ms-opentype" + }.get(path.suffix, mimetypes.guess_type(path)[0]) + if not mime_type: + raise ValueError("The attachment mime-type could not be automatically detected.") + + self.path = path + self.name = name + self.mime_type = mime_type + self.description = description + + def __repr__(self) -> str: + return "{name}({items})".format( + name=self.__class__.__name__, + items=", ".join([f"{k}={repr(v)}" for k, v in self.__dict__.items()]) + ) + + def __str__(self) -> str: + return " | ".join(filter(bool, [ + "ATT", + self.name, + self.mime_type, + self.description + ])) + + @property + def id(self) -> str: + """Compute an ID from the attachment data.""" + checksum = crc32(self.path.read_bytes()) + return hex(checksum) + + +__all__ = ("Attachment",) diff --git a/devine/core/tracks/tracks.py b/devine/core/tracks/tracks.py index e099dbb..4bdc7fc 100644 --- a/devine/core/tracks/tracks.py +++ b/devine/core/tracks/tracks.py @@ -14,6 +14,7 @@ from rich.tree import Tree from devine.core.config import config from devine.core.console import console from devine.core.constants import LANGUAGE_MAX_DISTANCE, AnyTrack, TrackT +from devine.core.tracks.attachment import Attachment from devine.core.tracks.audio import Audio from devine.core.tracks.chapters import Chapter, Chapters from devine.core.tracks.subtitle import Subtitle @@ -25,7 +26,7 @@ from devine.core.utils.collections import as_list, flatten class Tracks: """ - Video, Audio, Subtitle, and Chapter Track Store. + Video, Audio, Subtitle, Chapter, and Attachment Track Store. It provides convenience functions for listing, sorting, and selecting tracks. """ @@ -33,14 +34,23 @@ class Tracks: Video: 0, Audio: 1, Subtitle: 2, - Chapter: 3 + Chapter: 3, + Attachment: 4 } - def __init__(self, *args: Union[Tracks, Sequence[Union[AnyTrack, Chapter, Chapters]], Track, Chapter, Chapters]): + def __init__(self, *args: Union[ + Tracks, + Sequence[Union[AnyTrack, Chapter, Chapters, Attachment]], + Track, + Chapter, + Chapters, + Attachment + ]): self.videos: list[Video] = [] self.audio: list[Audio] = [] self.subtitles: list[Subtitle] = [] self.chapters = Chapters() + self.attachments: list[Attachment] = [] if args: self.add(args) @@ -53,7 +63,14 @@ class Tracks: def __add__( self, - other: Union[Tracks, Sequence[Union[AnyTrack, Chapter, Chapters]], Track, Chapter, Chapters] + other: Union[ + Tracks, + Sequence[Union[AnyTrack, Chapter, Chapters, Attachment]], + Track, + Chapter, + Chapters, + Attachment + ] ) -> Tracks: self.add(other) return self @@ -69,7 +86,8 @@ class Tracks: Video: [], Audio: [], Subtitle: [], - Chapter: [] + Chapter: [], + Attachment: [] } tracks = [*list(self), *self.chapters] @@ -98,7 +116,7 @@ class Tracks: return rep def tree(self, add_progress: bool = False) -> tuple[Tree, list[partial]]: - all_tracks = [*list(self), *self.chapters] + all_tracks = [*list(self), *self.chapters, *self.attachments] progress_callables = [] @@ -111,7 +129,7 @@ class Tracks: track_type_plural = track_type.__name__ + ("s" if track_type != Audio and num_tracks != 1 else "") tracks_tree = tree.add(f"[repr.number]{num_tracks}[/] {track_type_plural}") for track in tracks: - if add_progress and track_type != Chapter: + if add_progress and track_type not in (Chapter, Attachment): progress = Progress( SpinnerColumn(finished_text=""), BarColumn(), @@ -143,12 +161,19 @@ class Tracks: def add( self, - tracks: Union[Tracks, Sequence[Union[AnyTrack, Chapter, Chapters]], Track, Chapter, Chapters], + tracks: Union[ + Tracks, + Sequence[Union[AnyTrack, Chapter, Chapters, Attachment]], + Track, + Chapter, + Chapters, + Attachment + ], warn_only: bool = False ) -> None: """Add a provided track to its appropriate array and ensuring it's not a duplicate.""" if isinstance(tracks, Tracks): - tracks = [*list(tracks), *tracks.chapters] + tracks = [*list(tracks), *tracks.chapters, *tracks.attachments] duplicates = 0 for track in flatten(tracks): @@ -173,6 +198,8 @@ class Tracks: self.subtitles.append(track) elif isinstance(track, Chapter): self.chapters.add(track) + elif isinstance(track, Attachment): + self.attachments.append(track) else: raise ValueError("Track type was not set or is invalid.") @@ -363,6 +390,16 @@ class Tracks: else: chapters_path = None + for attachment in self.attachments: + if not attachment.path or not attachment.path.exists(): + raise ValueError("Attachment File was not found...") + cl.extend([ + "--attachment-description", attachment.description or "", + "--attachment-mime-type", attachment.mime_type, + "--attachment-name", attachment.name, + "--attach-file", str(attachment.path.resolve()) + ]) + output_path = ( self.videos[0].path.with_suffix(".muxed.mkv") if self.videos else self.audio[0].path.with_suffix(".muxed.mka") if self.audio else