mkvpriority 1.5.0__py3-none-any.whl

mkvpriority/__init__.py

@@ -0,0 +1,42 @@
+import importlib.metadata
+import tomllib
+from pathlib import Path
+
+from .main import (
+    Config,
+    Database,
+    Extension,
+    Track,
+    extract_tracks,
+    identify_tracks,
+    modify_tracks,
+    process_file,
+    process_tracks,
+    restore_file,
+    restore_tracks,
+)
+
+try:
+    __version__ = importlib.metadata.version('mkvpriority')
+except importlib.metadata.PackageNotFoundError:
+    __version__ = '(Unknown Version)'
+    try:
+        with Path('/app/pyproject.toml').open('rb') as f:
+            __version__ = tomllib.load(f)['project']['version']
+    except FileNotFoundError, KeyError:
+        pass
+
+
+__all__ = [
+    'Config',
+    'Database',
+    'Extension',
+    'Track',
+    'extract_tracks',
+    'identify_tracks',
+    'modify_tracks',
+    'process_file',
+    'process_tracks',
+    'restore_file',
+    'restore_tracks',
+]

mkvpriority/__main__.py

@@ -0,0 +1,4 @@
+from .main import main
+
+if __name__ == '__main__':
+    main()

mkvpriority/entrypoint.py

@@ -0,0 +1,201 @@
+import asyncio
+import logging
+import os
+import re
+import shlex
+import shutil
+import signal
+from pathlib import Path
+from typing import cast
+
+import cron_descriptor
+import pycountry
+from aiohttp import web
+from apscheduler.schedulers.asyncio import AsyncIOScheduler
+from apscheduler.triggers.cron import CronTrigger
+from cron_descriptor import FormatError
+
+import mkvpriority
+from mkvpriority import __version__
+from mkvpriority.main import setup_logging
+
+entrypoint_logger = logging.getLogger('entrypoint')
+processing_queue: asyncio.Queue[tuple[str, str, str | None]] = asyncio.Queue()
+
+
+MKVPRIORITY_ARGS = ['-c', '/config/config.toml'] + shlex.split(os.getenv('MKVPRIORITY_ARGS', ''))
+LOG_MAX_BYTES, LOG_MAX_FILES = os.getenv('LOG_MAX_BYTES'), os.getenv('LOG_MAX_FILES')
+
+CUSTOM_SCRIPT = os.getenv('CUSTOM_SCRIPT', 'false').lower() in ('true', '1', 't')
+WEBHOOK_RECEIVER = os.getenv('WEBHOOK_RECEIVER', 'false').lower() in ('true', '1', 't')
+WEBHOOK_PORT_STR = os.getenv('WEBHOOK_PORT') or ('8080' if WEBHOOK_RECEIVER else None)
+WEBHOOK_PORT = int(WEBHOOK_PORT_STR) if WEBHOOK_PORT_STR else None
+
+CRON_MACROS = {
+    '@yearly': '0 0 1 1 *',
+    '@annually': '0 0 1 1 *',
+    '@monthly': '0 0 1 * *',
+    '@weekly': '0 0 * * 0',
+    '@daily': '0 0 * * *',
+    '@midnight': '0 0 * * *',
+    '@hourly': '0 * * * *',
+}
+CRON_SCHEDULE = os.getenv('CRON_SCHEDULE')
+CRON_TARGET_PATHS = shlex.split(os.getenv('CRON_TARGET_PATHS', ''))
+
+
+def get_alpha_3_code(lang_name: str) -> str | None:
+    try:
+        lang = pycountry.languages.lookup(lang_name)
+        return cast(str, lang.alpha_3)  # ISO 639-3
+    except LookupError:
+        return None
+
+
+async def process_item(file_path: str, item_tags: str, orig_lang: str | None) -> None:
+    if item_tags:
+        file_path += f'::{re.split(r"[,;|]", item_tags)[0]}'
+    try:
+        argv = [*MKVPRIORITY_ARGS, file_path]
+        await asyncio.to_thread(mkvpriority.main.main, argv, orig_lang)
+    except Exception:
+        entrypoint_logger.exception(f"error occurred: '{file_path}'")
+
+
+async def queue_worker() -> None:
+    while True:
+        file_path, item_tags, item_id = await processing_queue.get()
+        await process_item(file_path, item_tags, item_id)
+        processing_queue.task_done()
+
+
+async def process_handler(request: web.Request) -> web.Response:
+    args = await request.json()
+    file_path = args.get('file_path')
+    item_tags = args.get('item_tags', '')
+    orig_lang = get_alpha_3_code(args.get('orig_lang', ''))
+    await processing_queue.put((file_path, item_tags, orig_lang))
+    return web.json_response({'message': f"received '{file_path}'"})
+
+
+async def create_runner(host: str, port: int) -> web.AppRunner:
+    app = web.Application()
+    app.router.add_post('/process', process_handler)
+
+    async def on_startup(app: web.Application) -> None:
+        app['worker'] = asyncio.create_task(queue_worker())
+
+    async def on_cleanup(app: web.Application) -> None:
+        app['worker'].cancel()
+        try:
+            await app['worker']
+        except asyncio.CancelledError:
+            pass
+
+    app.on_startup.append(on_startup)
+    app.on_cleanup.append(on_cleanup)
+
+    runner = web.AppRunner(app)
+    await runner.setup()
+    site = web.TCPSite(runner, host, port)
+    await site.start()
+    return runner
+
+
+async def create_scheduler(expr: str, timezone: str | None) -> AsyncIOScheduler:
+    scheduler = AsyncIOScheduler()
+    trigger = CronTrigger.from_crontab(expr, timezone)
+    cron_argv = MKVPRIORITY_ARGS + CRON_TARGET_PATHS
+    scheduler.add_job(lambda: mkvpriority.main.main(cron_argv), trigger)
+    scheduler.start()
+    return scheduler
+
+
+def main() -> None:
+    config_dir = Path('/config')
+    try:
+        config_dir.mkdir(parents=True, exist_ok=True)
+        config_file = config_dir / 'config.toml'
+        if not config_file.is_file():
+            shutil.copy2('config.toml', config_file)
+
+        extensions_dir = Path('/config/extensions')
+        extensions_dir.mkdir(parents=True, exist_ok=True)
+        init_file = extensions_dir / '__init__.py'
+        init_file.touch(exist_ok=True)
+
+        script_file = config_dir / 'mkvpriority.sh'
+        if not script_file.is_file():
+            shutil.copy2('mkvpriority.sh', script_file)
+
+        database_file = config_dir / 'archive.db'
+        database_file.touch(exist_ok=True)
+    except PermissionError:
+        entrypoint_logger.warning(f'recreate {config_dir} with correct PUID/PGID ownership')
+        raise
+
+    max_bytes = 5242880 if LOG_MAX_BYTES is None else int(LOG_MAX_BYTES)
+    max_files = 3 if LOG_MAX_FILES is None else int(LOG_MAX_FILES)
+    setup_logging('/config/mkvpriority.log', max_bytes, max_files)
+
+    entrypoint_logger.setLevel(logging.INFO)
+    logging.getLogger('aiohttp.access').setLevel(logging.WARNING)
+
+    async def run_all() -> None:
+        stop_event = asyncio.Event()
+        loop = asyncio.get_running_loop()
+
+        def handle_signal(sig_name: str) -> None:
+            entrypoint_logger.info(f'received {sig_name} signal')
+            stop_event.set()
+
+        loop.add_signal_handler(signal.SIGTERM, lambda: handle_signal('SIGTERM'))
+        loop.add_signal_handler(signal.SIGINT, lambda: handle_signal('SIGINT'))
+
+        try:
+            runner = None
+            if WEBHOOK_PORT:
+                runner = await create_runner('0.0.0.0', WEBHOOK_PORT)
+                entrypoint_logger.info(f'webhook listener started on 0.0.0.0:{WEBHOOK_PORT}')
+
+            scheduler = None
+            if expr := CRON_SCHEDULE:
+                if expr.startswith('@'):
+                    macro = expr
+                    try:
+                        expr = CRON_MACROS[macro]
+                    except KeyError as e:
+                        e.add_note(f"unsupported cron macro '{macro}'")
+                        raise
+                timezone = os.getenv('TZ', 'UTC')
+                try:
+                    expr_desc = cron_descriptor.get_description(expr)
+                    expr_desc = expr_desc[0].lower() + expr_desc[1:]
+                    scheduler = await create_scheduler(expr, timezone)
+                except (FormatError, ValueError) as e:
+                    e.add_note(f"unsupported cron expression '{expr}'")
+                    raise
+                entrypoint_logger.info(f'scheduled task to run {expr_desc} ({timezone})')
+
+            await stop_event.wait()
+
+        finally:
+            if scheduler:
+                scheduler.shutdown(wait=False)
+            if runner:
+                await runner.cleanup()
+
+    entrypoint_logger.info(f'MKVPriority {__version__}')
+    asyncio.run(run_all())
+
+
+if __name__ == '__main__':
+    if CUSTOM_SCRIPT or WEBHOOK_RECEIVER:
+        entrypoint_logger.warning(
+            'CUSTOM_SCRIPT and WEBHOOK_RECEIVER are deprecated; use WEBHOOK_PORT instead'
+        )
+        main()
+    elif WEBHOOK_PORT or CRON_SCHEDULE:
+        main()
+    else:
+        mkvpriority.main.main()

mkvpriority/extensions/multiplexer.py

@@ -0,0 +1,112 @@
+import json
+import subprocess
+import tomllib
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Any
+
+from mkvpriority import Config, Extension, Track
+from mkvpriority.main import mkvmerge_logger
+
+
+class Multiplexer(Extension):
+    def __init__(self) -> None:
+        super().__init__('multiplexer')
+        self.parameters: dict[str, Any] = {}
+
+    def process_file(
+        self,
+        file_path: Path,
+        video_tracks: list[Track],
+        audio_tracks: list[Track],
+        subtitle_tracks: list[Track],
+        config: Config,
+        dry_run: bool = False,
+    ) -> None:
+        if config.toml_path in self.parameters:
+            parameters = self.parameters[config.toml_path]
+        else:
+            with open(config.toml_path, 'rb') as f:
+                toml_file = tomllib.load(f)
+            parameters = toml_file.get('multiplexer', {})
+            self.parameters[config.toml_path] = parameters
+        self.strip: bool = parameters.get('strip_tracks', False)
+        self.reorder: bool = parameters.get('reorder_tracks', False)
+        self.filter_tracks(file_path, video_tracks, audio_tracks, subtitle_tracks, config, dry_run)
+
+    def multiplex_tracks(self, arguments: list[str]) -> None:
+        with NamedTemporaryFile('w+', encoding='utf-8', suffix='.json', delete=False) as temp_file:
+            json.dump(arguments, temp_file)
+            temp_file_path = Path(temp_file.name)
+
+        try:
+            result = subprocess.run(
+                ['mkvmerge', f'@{temp_file_path}'],
+                capture_output=True,
+                encoding='utf-8',
+                check=True,
+                text=True,
+            )
+            mkvmerge_logger.debug(result.stdout.strip())
+        finally:
+            temp_file_path.unlink(missing_ok=True)
+
+    def filter_tracks(
+        self,
+        file_path: Path,
+        video_tracks: list[Track],
+        audio_tracks: list[Track],
+        subtitle_tracks: list[Track],
+        config: Config,
+        dry_run: bool = False,
+    ) -> None:
+        track_order: list[str] = []
+        audio_strip: list[str] = []
+        subtitle_strip: list[str] = []
+
+        def process_tracks(
+            tracks: list[Track],
+            track_order: list[str],
+            track_strip: list[str],
+            track_langs: dict[str, int],
+        ) -> None:
+            for track in tracks:
+                if self.strip and self.reorder:
+                    if track.language in track_langs:
+                        track_order.append(f'0:{track.index}')
+                    else:
+                        track_strip.append(f'!{track.index}')
+                elif self.strip:
+                    if track.language not in track_langs:
+                        track_strip.append(f'!{track.index}')
+                elif self.reorder:
+                    track_order.append(f'0:{track.index}')
+
+        for track in video_tracks:
+            track_order.append(f'0:{track.index}')
+        process_tracks(audio_tracks, track_order, audio_strip, config.audio_languages)
+        process_tracks(subtitle_tracks, track_order, subtitle_strip, config.subtitle_languages)
+
+        temp_output_path = file_path.with_name(f'{file_path.stem}_temp.mkv')
+        mkv_args = ['-o', str(temp_output_path)]
+        if audio_strip:
+            mkv_args += ['--audio-tracks', ','.join(audio_strip)]
+        if subtitle_strip:
+            mkv_args += ['--subtitle-tracks', ','.join(subtitle_strip)]
+        mkv_args += [str(file_path)]
+        if track_order and any(
+            int(id_a.split(':')[1]) > int(id_b.split(':')[1])
+            for id_a, id_b in zip(track_order, track_order[1:])
+        ):
+            mkv_args += ['--track-order', ','.join(track_order)]
+
+        if len(mkv_args) > 3:
+            self.extension_logger.info(' '.join(mkv_args))
+            if not dry_run:
+                try:
+                    self.multiplex_tracks(mkv_args)
+                    temp_output_path.replace(file_path)
+                except subprocess.CalledProcessError as e:
+                    mkvmerge_logger.error((e.stderr or e.stdout or str(e)).strip())
+                    temp_output_path.unlink(missing_ok=True)
+                    raise

mkvpriority/extensions/subtitle_extractor.py

@@ -0,0 +1,74 @@
+import json
+import subprocess
+import tomllib
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Any
+
+from mkvpriority import Config, Extension, Track
+
+SUBTITLE_EXTENSIONS = {'ASS': 'ass', 'SSA': 'ssa', 'UTF8': 'srt', 'WEBVTT': 'vtt'}
+
+
+class SubtitleExtractor(Extension):
+    def __init__(self) -> None:
+        super().__init__('subtitle_extractor')
+        self.parameters: dict[str, Any] = {}
+
+    def process_file(
+        self,
+        file_path: Path,
+        video_tracks: list[Track],
+        audio_tracks: list[Track],
+        subtitle_tracks: list[Track],
+        config: Config,
+        dry_run: bool = False,
+    ) -> None:
+        if not subtitle_tracks:
+            return
+
+        if config.toml_path in self.parameters:
+            extract = self.parameters[config.toml_path]
+        else:
+            with open(config.toml_path, 'rb') as f:
+                toml_file = tomllib.load(f)
+            extract = toml_file.get('extract_embedded_subtitles', False)
+            self.parameters[config.toml_path] = extract
+
+        if extract:
+            subtitle_track = max(subtitle_tracks, key=lambda track: track.score)
+            subtitle_path = self.build_subtitle_path(file_path, subtitle_track)
+            if subtitle_path and not subtitle_path.is_file():
+                self.extract_subtitles(file_path, subtitle_path, subtitle_track.index)
+
+    def build_subtitle_path(self, file_path: Path, subtitle_track: Track) -> Path | None:
+        if not subtitle_track.codec.startswith('S_TEXT/'):
+            return None
+        subtitle_format = subtitle_track.codec.split('/')[-1]
+        if subtitle_format not in SUBTITLE_EXTENSIONS:
+            return None
+        extension = SUBTITLE_EXTENSIONS[subtitle_format]
+        subtitle_suffix = f'.{subtitle_track.language}'
+        if subtitle_track.default:
+            subtitle_suffix += '.default'
+        if subtitle_track.forced:
+            subtitle_suffix += '.forced'
+        return Path(file_path).with_suffix(f'{subtitle_suffix}.{extension}')
+
+    def extract_subtitles(self, file_path: Path, subtitle_path: Path, index: int) -> None:
+        self.extension_logger.info(f"extracting embedded subtitles to '{subtitle_path}'")
+        with NamedTemporaryFile('w+', encoding='utf-8', suffix='.json', delete=False) as temp_file:
+            json.dump(['tracks', str(file_path), f'{index}:{subtitle_path}'], temp_file)
+            temp_file_path = Path(temp_file.name)
+
+        try:
+            result = subprocess.run(
+                ['mkvextract', f'@{temp_file_path}'],
+                capture_output=True,
+                encoding='utf-8',
+                check=True,
+                text=True,
+            )
+            self.extension_logger.debug(result.stdout.strip())
+        finally:
+            temp_file_path.unlink(missing_ok=True)

mkvpriority/extensions/subtitle_restyler.py

@@ -0,0 +1,199 @@
+import re
+import tomllib
+from collections import defaultdict
+from pathlib import Path
+from typing import Any
+
+from mkvpriority import Config, Extension, Track
+
+SAFE_ATTRS = {
+    'Fontname',
+    'PrimaryColour',
+    'SecondaryColour',
+    'OutlineColour',
+    'BackColour',
+    'Bold',
+    'Italic',
+    'Underline',
+    'StrikeOut',
+    'ScaleX',
+    'ScaleY',
+    'Angle',
+    'Alignment',
+    'BorderStyle',
+    'Encoding',
+}
+RES_DEP_X = {'Spacing', 'MarginL', 'MarginR'}
+RES_DEP_Y = {'Fontsize', 'Outline', 'Shadow', 'MarginV'}
+ASS_ATTR_MAP = {attr.lower(): attr for attr in SAFE_ATTRS | RES_DEP_X | RES_DEP_Y}
+
+
+class SubtitleRestyler(Extension):
+    def __init__(self, max_ratio: float = 0.15, max_allowance: int = 2):
+        super().__init__('subtitle_restyler')
+        self.parameters: dict[str, Any] = {}
+        self.max_ratio = max_ratio
+        self.max_allowance = max_allowance
+
+    def process_file(
+        self,
+        file_path: Path,
+        video_tracks: list[Track],
+        audio_tracks: list[Track],
+        subtitle_tracks: list[Track],
+        config: Config,
+        dry_run: bool = False,
+    ) -> None:
+        if not subtitle_tracks:
+            return
+
+        if config.toml_path in self.parameters:
+            attributes = self.parameters[config.toml_path]
+        else:
+            with open(config.toml_path, 'rb') as f:
+                toml_file = tomllib.load(f)
+            attributes = toml_file.get('subtitle_styles', {})
+            self.parameters[config.toml_path] = attributes
+
+        if attributes:
+            subtitle_track = max(subtitle_tracks, key=lambda track: track.score)
+            subtitle_path = self.build_subtitle_path(file_path, subtitle_track)
+            if subtitle_path and subtitle_path.is_file():
+                self.modify_subtitle_styles(subtitle_path, attributes)
+
+    def build_subtitle_path(self, file_path: Path, subtitle_track: Track) -> Path | None:
+        if not subtitle_track.codec.startswith('S_TEXT/'):
+            return None
+        if subtitle_track.codec.split('/')[-1] != 'ASS':
+            return None
+        subtitle_suffix = f'.{subtitle_track.language}'
+        if subtitle_track.default:
+            subtitle_suffix += '.default'
+        if subtitle_track.forced:
+            subtitle_suffix += '.forced'
+        return Path(file_path).with_suffix(f'{subtitle_suffix}.ass')
+
+    def scale_style_attributes(
+        self, input_lines: list[str], attributes: dict[str, Any]
+    ) -> dict[str, str]:
+        playres_x, playres_y = 1920.0, 1080.0
+        for line in input_lines:
+            if line.startswith('PlayResX:'):
+                playres_x = float(line.split(':')[1].strip())
+            elif line.startswith('PlayResY:'):
+                playres_y = float(line.split(':')[1].strip())
+            elif line.startswith('[Events]'):
+                break
+
+        scale_x = playres_x / 1920.0
+        scale_y = playres_y / 1080.0
+        scaled_attributes: dict[str, str] = {}
+        for attr, val in attributes.items():
+            attr = ASS_ATTR_MAP.get(attr.lower(), attr)
+            if attr in SAFE_ATTRS:
+                scaled_attributes[attr] = str(val)
+            elif attr in RES_DEP_X:
+                scaled_val = float(val) * scale_x
+                scaled_attributes[attr] = str(
+                    int(round(scaled_val)) if 'Margin' in attr else round(scaled_val, 2)
+                )
+            elif attr in RES_DEP_Y:
+                scaled_val = float(val) * scale_y
+                scaled_attributes[attr] = str(
+                    int(round(scaled_val)) if 'Margin' in attr else round(scaled_val, 2)
+                )
+            else:
+                self.extension_logger.warning(f"style '{attr}' is not in [V4+ Styles]")
+
+        return scaled_attributes
+
+    def detect_dialogue_styles(self, input_lines: list[str]) -> set[str]:
+        style_stats: dict[str, dict[str, int]] = defaultdict(
+            lambda: {'count_spatial': 0, 'count_karaoke': 0, 'count_drawing': 0, 'total': 0}
+        )
+        in_events_section = False
+
+        position_pattern = re.compile(r'\\(pos|move|org|clip|iclip)\s*\(|\\an[1-9]', re.IGNORECASE)
+        rotation_pattern = re.compile(r'\\(fr[xyz]|fa[xy])\s*-?\d', re.IGNORECASE)
+        karaoke_pattern = re.compile(r'\\[kK][fo]?[0-9]+')
+        drawing_pattern = re.compile(r'\\[pP][1-9]')
+
+        for line in input_lines:
+            if line.startswith('[Events]'):
+                in_events_section = True
+                continue
+            elif line.startswith('['):
+                in_events_section = False
+
+            if in_events_section and line.startswith('Dialogue:'):
+                parts = line.split(':', 1)[1].strip().split(',', 9)
+                if len(parts) > 9:
+                    style_name = parts[3].strip()
+                    text = parts[9]
+                    style_stats[style_name]['total'] += 1
+                    if position_pattern.search(text) or rotation_pattern.search(text):
+                        style_stats[style_name]['count_spatial'] += 1
+                    if karaoke_pattern.search(text):
+                        style_stats[style_name]['count_karaoke'] += 1
+                    if drawing_pattern.search(text):
+                        style_stats[style_name]['count_drawing'] += 1
+
+        subtitle_styles: set[str] = set()
+        for style, stats in style_stats.items():
+            ratio_karaoke = stats['count_karaoke'] / stats['total']
+            if stats['count_karaoke'] > 0 and ratio_karaoke > 0.1:
+                continue
+            if stats['count_drawing'] > 0:
+                continue
+            ratio_spatial = stats['count_spatial'] / stats['total']
+            is_dialogue = ratio_spatial <= self.max_ratio
+            if not is_dialogue and stats['count_spatial'] <= self.max_allowance:
+                if ratio_spatial < 1.0:
+                    is_dialogue = True
+            if is_dialogue:
+                subtitle_styles.add(style)
+
+        return subtitle_styles
+
+    def modify_subtitle_styles(self, file_path: Path, attributes: dict[str, Any]) -> None:
+        with open(file_path, encoding='utf-8-sig') as f:
+            input_lines = f.readlines()
+        scaled_attributes = self.scale_style_attributes(input_lines, attributes)
+        if not scaled_attributes:
+            return
+        subtitle_styles = self.detect_dialogue_styles(input_lines)
+        if not subtitle_styles:
+            return
+
+        attr_indices: dict[str, int] = {}
+        output_lines: list[str] = []
+        in_styles_section = False
+        for line in input_lines:
+            if line.startswith('[V4+ Styles]'):
+                in_styles_section = True
+                output_lines.append(line)
+                continue
+            elif line.startswith('['):
+                in_styles_section = False
+
+            if in_styles_section:
+                if line.startswith('Format:'):
+                    format_string = line.split(':', 1)[1].strip()
+                    format_parts = [p.strip() for p in format_string.split(',')]
+                    for attr in scaled_attributes.keys():
+                        if attr in format_parts:
+                            attr_indices[attr] = format_parts.index(attr)
+                elif line.startswith('Style:') and attr_indices:
+                    style_parts = line.split(':', 1)[1].strip().split(',')
+                    style_name = style_parts[0].strip()
+                    if style_name in subtitle_styles:
+                        for attr, val in scaled_attributes.items():
+                            if attr in attr_indices:
+                                index = attr_indices[attr]
+                                style_parts[index] = val
+                        line = 'Style: ' + ','.join(style_parts) + '\n'
+            output_lines.append(line)
+
+        self.extension_logger.info(f'restyling external subtitles for {sorted(subtitle_styles)}')
+        with open(file_path, 'w', encoding='utf-8-sig') as f:
+            f.writelines(output_lines)