Add new flag --use-ass-in-mkv

Author: moi15moi

README.md

@@ -17,7 +17,7 @@ pip install FontCollector
 ## FontCollector Usage
 ```console
 $ fontcollector --help
-usage: fontcollector [-h] --input INPUT [INPUT ...] [-mkv MKV] [--output OUTPUT] [-mkvtoolnix MKVTOOLNIX] [--delete-fonts] [--additional-fonts ADDITIONAL_FONTS [ADDITIONAL_FONTS ...]]
+usage: fontcollector [-h] [--input INPUT [INPUT ...]] [-mkv MKV] [--use-ass-in-mkv] [--output OUTPUT] [-mkvtoolnix MKVTOOLNIX] [--delete-fonts] [--additional-fonts ADDITIONAL_FONTS [ADDITIONAL_FONTS ...]]
                      [--additional-fonts-recursive ADDITIONAL_FONTS_RECURSIVE [ADDITIONAL_FONTS_RECURSIVE ...]] [--exclude-system-fonts] [--collect-draw-fonts] [--dont-convert-variable-to-collection] [--logging [LOGGING]]
 
 FontCollector for Advanced SubStation Alpha file.
@@ -28,6 +28,8 @@ options:
                         Subtitles file. Must be an ASS file/directory. You can specify more than one .ass file/path.
   -mkv MKV
                         Video where the fonts will be merge. Must be a Matroska file.
+  --use-ass-in-mkv, -ass-mkv
+                        If specified, it will use the .ass file muxed to the mkv and collect those fonts and mux them to the mkv. If not specified, it will do nothing.
   --output OUTPUT, -o OUTPUT
                         Destination path of the font. If -o and -mkv aren't specified, it will be the current path.
   -mkvtoolnix MKVTOOLNIX

font_collector/__main__.py

@@ -2,6 +2,7 @@
 import shutil
 from pathlib import Path
 from sys import argv
+from tempfile import TemporaryDirectory
 
 from . import _handler
 from .ass.ass_document import AssDocument
@@ -12,7 +13,7 @@
     FontLoader,
     FontSelectionStrategyLibass
 )
-from .mkvtoolnix.mkvpropedit import MKVPropedit
+from .mkvtoolnix import MKVPropedit, MKVExtract
 from .parse_arguments import parse_arguments
 
 _logger = logging.getLogger(__name__)
@@ -23,6 +24,7 @@ def main() -> None:
         ass_files_path,
         output_directory,
         mkv_path,
+        use_ass_in_mkv,
         delete_fonts,
         additional_fonts_path,
         additional_fonts_recursive_path,
@@ -52,6 +54,20 @@ def main() -> None:
 
             fonts_file_found.update(collect_subtitle_fonts(subtitle, font_collection, font_strategy, collect_draw_fonts, convert_variable_to_collection, output_directory))
             _logger.info("")
+        
+        if use_ass_in_mkv:
+            with TemporaryDirectory() as tmp_dir:
+                assert isinstance(mkv_path, Path)
+                mkv_ass_files = MKVExtract.get_mkv_ass_files(mkv_path, Path(tmp_dir))
+                for mkv_ass_file in mkv_ass_files:
+                    subtitle = AssDocument.from_file(mkv_ass_file.filename)
+                    log_msg = f"Loaded successfully the .ass stream at index {mkv_ass_file.mkv_id}"
+                    if mkv_ass_file.track_name:
+                        log_msg += f" - \"{mkv_ass_file.track_name}\""
+                    _logger.info(log_msg)
+
+                    fonts_file_found.update(collect_subtitle_fonts(subtitle, font_collection, font_strategy, collect_draw_fonts, convert_variable_to_collection, output_directory))
+                    _logger.info("")
 
         if mkv_path is not None:
             if delete_fonts:

font_collector/parse_arguments.py

@@ -7,7 +7,9 @@
 from .mkvtoolnix.mkv_utils import MKVUtils
 
 
-def __parse_input_file(ass_input: list[Path]) -> list[Path]:
+def __parse_input_file(ass_input: Optional[list[Path]]) -> list[Path]:
+    if ass_input is None:
+        return []
 
     ass_files_path = []
     for input in ass_input:
@@ -33,6 +35,7 @@ def parse_arguments() -> tuple[
     Path,
     Union[Path, None],
     bool,
+    bool,
     Iterable[Path],
     Iterable[Path],
     bool,
@@ -56,7 +59,6 @@ def parse_arguments() -> tuple[
         "-i",
         nargs="+",
         type=Path,
-        required=True,
         help="""
     Subtitles file. Must be an ASS file/directory. You can specify more than one .ass file/path.
     """,
@@ -68,6 +70,14 @@ def parse_arguments() -> tuple[
     Video where the fonts will be merge. Must be a Matroska file.
     """,
     )
+    parser.add_argument(
+        "--use-ass-in-mkv",
+        "-ass-mkv",
+        action="store_true",
+        help="""
+    If specified, it will use the .ass file muxed to the mkv and collect those fonts and mux them to the mkv. If not specified, it will do nothing.
+    """,
+    )
     parser.add_argument(
         "--output",
         "-o",
@@ -150,12 +160,9 @@ def parse_arguments() -> tuple[
 
     # Parse args
     ass_files_path = __parse_input_file(args.input)
-
-    if len(ass_files_path) == 0:
-        raise RuntimeError("The specified file(s)/folder(s) doesn't exist or the folder(s) doesn't contains any .ass file.")
-
     output_directory = args.output
     mkv_path = args.mkv
+    use_ass_in_mkv = args.use_ass_in_mkv
     delete_fonts = args.delete_fonts
     additional_fonts = args.additional_fonts
     additional_fonts_recursive = args.additional_fonts_recursive
@@ -164,6 +171,12 @@ def parse_arguments() -> tuple[
     convert_variable_to_collection = args.dont_convert_variable_to_collection
     logging_file_path = args.logging
 
+    if len(ass_files_path) == 0 and not use_ass_in_mkv:
+        raise RuntimeError("The specified file(s)/folder(s) doesn't exist or the folder(s) doesn't contains any .ass file.")
+    
+    if use_ass_in_mkv and mkv_path is None:
+        raise RuntimeError("You need to add the flag `-mkv` to use the flag `--use-ass-in-mkv`.")
+
     if args.mkvtoolnix:
         if not mkv_path:
             raise RuntimeError("-mkvtoolnix requires --mkv option.")
@@ -173,6 +186,7 @@ def parse_arguments() -> tuple[
         ass_files_path,
         output_directory,
         mkv_path,
+        use_ass_in_mkv,
         delete_fonts,
         additional_fonts,
         additional_fonts_recursive,