diff --git a/pymkv/MKVAttachment.py b/pymkv/MKVAttachment.py index 18a46ba..3e72151 100644 --- a/pymkv/MKVAttachment.py +++ b/pymkv/MKVAttachment.py @@ -55,10 +55,11 @@ class MKVAttachment: which will attach to all files. """ - def __init__(self, file_path, name=None, description=None, attach_once=False): + def __init__(self, file_path, attachment_id=1, name=None, description=None, attach_once=False): self.mime_type = None self._file_path = None self.file_path = file_path + self.attachment_id = attachment_id self.name = name self.description = description self.attach_once = attach_once diff --git a/pymkv/MKVFile.py b/pymkv/MKVFile.py index afc234c..4ffaed3 100644 --- a/pymkv/MKVFile.py +++ b/pymkv/MKVFile.py @@ -36,8 +36,8 @@ """ import json -from os import devnull -from os.path import expanduser, isfile +from os import devnull, makedirs +from os.path import expanduser, isfile, isdir import subprocess as sp import bitmath @@ -46,7 +46,7 @@ from pymkv.MKVAttachment import MKVAttachment from pymkv.Timestamp import Timestamp from pymkv.ISO639_2 import is_ISO639_2 -from pymkv.Verifications import verify_matroska, verify_mkvmerge +from pymkv.Verifications import verify_matroska, verify_mkvmerge, verify_mkvextract class MKVFile: @@ -81,6 +81,8 @@ class MKVFile: def __init__(self, file_path=None, title=None): self.mkvmerge_path = 'mkvmerge' + self.mkvextract_path = 'mkvextract' + self.file_path = file_path self.title = title self._chapters_file = None self._chapter_language = None @@ -111,6 +113,12 @@ def __init__(self, file_path=None, title=None): if 'forced_track' in track['properties']: new_track.forced_track = track['properties']['forced_track'] self.add_track(new_track) + + # add attachments with info + for attachment in info_json['attachments']: + new_attachment = MKVAttachment(file_path, attachment['id'], attachment['file_name'], attachment['description']) + new_attachment.mime_type = attachment['content_type'] + self.add_attachment(new_attachment) # split options self._split_options = [] @@ -189,20 +197,23 @@ def command(self, output_path, subprocess=False): command.extend(['-s', str(track.track_id)]) # exclusions + command.append('--no-attachments') if track.no_chapters: command.append('--no-chapters') if track.no_global_tags: command.append('--no-global-tags') if track.no_track_tags: command.append('--no-track-tags') - if track.no_attachments: - command.append('--no-attachments') # add path command.append(track.file_path) # add attachments - for attachment in self.attachments: + own_attachments = [a for a in self.attachments if a.file_path == self.file_path] + if own_attachments: + command += ['-D', '-A', '-S', '-B', '-T', '--no-chapters', '-m', ",".join(str(a.attachment_id) for a in own_attachments), self.file_path] + + for attachment in (a for a in self.attachments if a.file_path != self.file_path): # info if attachment.name is not None: command.extend(['--attachment-name', attachment.name]) @@ -475,6 +486,72 @@ def remove_track(self, track_num): else: raise IndexError('track index out of range') + def extract_track(self, track_num, output_path, silent=False): + """Extract a :class:`~pymkv.MKVTrack` from the :class:`~pymkv.MKVFile` object. + + Parameters + ---------- + track_num : int + Index of the track to extract. + output_path : str + The path to be used as the output file in the mkvextract command. + silent : bool, optional + By default the mkvextract output will be shown unless silent is True. + + Raises + ------ + FileNotFoundError + Raised if the path to mkvextract could not be verified. + IndexError + Raised if `track_num` is is out of range of the track list. + """ + if not verify_mkvextract(mkvextract_path=self.mkvextract_path): + raise FileNotFoundError('mkvextract is not at the specified path, add it there or change the mkvextract_path ' + 'property') + output_path = expanduser(output_path) + track = self.get_track(track_num) + command = [self.mkvextract_path, track.file_path, 'tracks', f'{track.track_id}:{output_path}'] + if silent: + sp.run(command, check=True, stdout=sp.DEVNULL, stderr=sp.DEVNULL) + else: + print(f'Running with command:\n"{command}"') + sp.run(command, check=True, capture_output=True) + + def extract_attachments(self, output_directory, attachment_ids=[], silent=False): + """Extract all :class:`~pymkv.MKVAttachment` from the :class:`~pymkv.MKVFile` object. + + Parameters + ---------- + output_directory : str + The output directory to be used in the mkvextract command. + attachment_ids : list[int] + The IDs of attachments to extract. + silent : bool, optional + By default the mkvextract output will be shown unless silent is True. + + Raises + ------ + FileNotFoundError + Raised if the path to mkvextract could not be verified. + """ + if not verify_mkvextract(mkvextract_path=self.mkvextract_path): + raise FileNotFoundError('mkvextract is not at the specified path, add it there or change the mkvextract_path ' + 'property') + output_directory = expanduser(output_directory) + if not isdir(output_directory): + makedirs(output_directory) + command = [self.mkvextract_path, self.file_path, 'attachments'] + own_attachments = [a for a in self.attachments if a.file_path == self.file_path] + if attachment_ids: + own_attachments = [a for a in own_attachments if a.attachment_id in attachment_ids] + for a in own_attachments: + command.append(f"{a.attachment_id}:{output_directory}/{a.name}") + if silent: + sp.run(command, check=True, stdout=sp.DEVNULL, stderr=sp.DEVNULL) + else: + print(f'Running with command:\n"{command}"') + sp.run(command, check=True, capture_output=True) + def split_none(self): """Remove all splitting options.""" self._split_options = [] diff --git a/pymkv/MKVTrack.py b/pymkv/MKVTrack.py index 329d4d1..2e9ecae 100644 --- a/pymkv/MKVTrack.py +++ b/pymkv/MKVTrack.py @@ -102,6 +102,7 @@ def __init__(self, file_path, track_id=0, track_name=None, language=None, defaul # track info self._track_codec = None self._track_type = None + self._audio_channels = 0 # base self.mkvmerge_path = 'mkvmerge' @@ -171,6 +172,8 @@ def track_id(self, track_id): self._track_id = track_id self._track_codec = info_json['tracks'][track_id]['codec'] self._track_type = info_json['tracks'][track_id]['type'] + if self._track_type == 'audio': + self._audio_channels = info_json['tracks'][track_id]['properties']['audio_channels'] @property def language(self): @@ -226,3 +229,8 @@ def track_codec(self): def track_type(self): """str: The type of track such as video or audio.""" return self._track_type + + @property + def audio_channels(self): + """int: The number of audio channels in the track.""" + return self._audio_channels diff --git a/pymkv/Verifications.py b/pymkv/Verifications.py index ad74270..00962d7 100644 --- a/pymkv/Verifications.py +++ b/pymkv/Verifications.py @@ -24,6 +24,19 @@ def verify_mkvmerge(mkvmerge_path='mkvmerge'): return True return False +def verify_mkvextract(mkvextract_path='mkvextract'): + """Verify mkvextract is working. + + mkvextract_path (str): + Alternate path to mkvextract if it is not already in the $PATH variable. + """ + try: + output = sp.check_output([mkvextract_path, '-V']).decode() + except (sp.CalledProcessError, FileNotFoundError): + return False + if match('mkvextract.*', output): + return True + return False def verify_matroska(file_path, mkvmerge_path='mkvmerge'): """Verify if a file is a Matroska file.