yt_common: colossally stupid ordered chapters stuff

.gitignore

@@ -25,4 +25,5 @@ ffmpeg2pass*.log
 **/bdmv/
 **/bdpremux/
 **/sub/
+**/dvd/
 *.txt

yt_common/.flake8

@@ -0,0 +1,7 @@
+[flake8]
+count = True
+ignore = W503,E226
+max-line-length = 120
+exclude = stubs/*
+show-source = True
+statistics = True

yt_common/yt_common/automation.py

@@ -7,19 +7,24 @@ import random
 import shutil
 import string
 import subprocess
+import tempfile
 
 from lvsfunc.render import clip_async_render
 
 from typing import Any, BinaryIO, Callable, List, Optional, Sequence, Set, Tuple, cast
 
-from .chapters import Chapter, make_chapters
+from .chapters import Chapter, Edition, make_chapters, make_qpfile
 from .config import Config
 from .logging import log
 from .source import FileSource
 
 core = vs.core
 
-AUDIO_ENCODE: str = "_audiogetter_encode.mka"
+AUDIO_PFX: str = "_audiogetter_temp_"
+
+
+def get_temp_filename(prefix: str = "", suffix: str = "") -> str:
+    return f"{prefix}{next(tempfile._get_candidate_names())}{suffix}"  # type: ignore
 
 
 def bin_to_plat(binary: str) -> str:
@@ -63,7 +68,7 @@ class Encoder():
             log.warn("Existing output detected, skipping encode!")
             return outfile, []
 
-        params = [p.format(frames=end-start, filename=filename) for p in self.params]
+        params = [p.format(frames=end-start, filename=filename, qpfile="qpfile.txt") for p in self.params]
 
         log.status("--- RUNNING ENCODE ---")
 
@@ -124,14 +129,24 @@ class AudioGetter():
         self.cleanup = set()
 
     def trim_audio(self, ftrim: Optional[acsuite.types.Trim] = None) -> str:
-        trims = self.src.get_audio()
-        if not trims or len(trims) > 1:
-            raise NotImplementedError("Please implement multifile trimming")
-        audio_cut = acsuite.eztrim(trims[0].path, trims[0].trim or (0, None), streams=0)[0]
-        self.cleanup.add(audio_cut)
+        trims = self.src.audio_src()
+
+        tlist: List[str] = []
+        for t in trims:
+            audio_cut = acsuite.eztrim(t.path, t.trim or (0, None), ref_clip=self.src.audio_ref(),
+                                       outfile=get_temp_filename(prefix=AUDIO_PFX, suffix=".mka"),
+                                       streams=0)[0]
+            self.cleanup.add(audio_cut)
+            tlist.append(audio_cut)
+
+        if len(tlist) > 1:
+            ffmpeg = acsuite.ffmpeg.FFmpegAudio()
+            audio_cut = ffmpeg.concat(*tlist)
+            self.cleanup.add(audio_cut)
 
         if ftrim:
-            audio_cut = acsuite.eztrim(audio_cut, ftrim, ref_clip=self.src.source())[0]
+            audio_cut = acsuite.eztrim(audio_cut, ftrim, ref_clip=self.src.source(),
+                                       outfile=get_temp_filename(prefix=AUDIO_PFX, suffix=".mka"))[0]
             self.cleanup.add(audio_cut)
 
         audio_cut = self.config.encode_audio(audio_cut)
@@ -164,7 +179,8 @@ class SelfRunner():
 
     def __init__(self, config: Config, source: FileSource, final_filter: Callable[[], vs.VideoNode],
                  workraw_filter: Optional[Callable[[], vs.VideoNode]] = None,
-                 chapters: Optional[List[Chapter]] = None) -> None:
+                 chapters: Optional[List[Chapter]] = None,
+                 editions: Optional[List[Edition]] = None) -> None:
         self.config = config
         self.src = source
         self.video_clean = False
@@ -242,6 +258,10 @@ class SelfRunner():
             log.status("Comparison generated.")
             return
 
+        open("qpfile.txt", "w").close()  # guarantee qpfile is created/cleared as appropriate
+        if (editions or chapters) and (start == 0 and end == self.clip.num_frames):
+            make_qpfile("qpfile.txt", chapters=chapters, editions=editions)
+
         settings_path = os.path.join(self.config.datapath, f"{self.profile}-settings")
         if not os.path.isfile(settings_path):
             raise FileNotFoundError(f"Failed to find {settings_path}!")
@@ -260,9 +280,9 @@ class SelfRunner():
 
         self._do_audio(start, audio_end)
 
-        if chapters:
+        if (editions or chapters) and (start == 0 and end == self.clip.num_frames):
             log.status("--- GENERATING CHAPTERS  ---")
-            make_chapters(chapters, self.timecodes, f"{self.config.desc}.xml")
+            make_chapters(self.timecodes, f"{self.config.desc}.xml", chapters=chapters, editions=editions)
 
         try:
             log.status("--- MUXING FILE ---")

yt_common/yt_common/chapters.py

@@ -1,7 +1,9 @@
 from lxml import etree
+from lxml.etree import _Element
+
 from random import getrandbits
 
-from typing import Dict, List, NamedTuple, Set
+from typing import Dict, List, NamedTuple, Optional, Set
 
 LANGMAP: Dict[str, str] = {
     "eng": "en",
@@ -13,9 +15,16 @@ LANGMAP: Dict[str, str] = {
 class Chapter(NamedTuple):
     title: str
     frame: int
+    end_frame: Optional[int] = None
     lang: str = "eng"
 
 
+class Edition(NamedTuple):
+    chapters: List[Chapter]
+    default: bool = False
+    ordered: bool = False
+
+
 class RandMan:
     used: Set[int]
 
@@ -30,6 +39,9 @@ class RandMan:
         return str(r)
 
 
+rand = RandMan()
+
+
 def timecode_to_timestamp(stamp: float) -> str:
     m = int(stamp // 60)
     stamp %= 60
@@ -38,18 +50,15 @@ def timecode_to_timestamp(stamp: float) -> str:
     return f"{h:02d}:{m:02d}:{stamp:06.3f}000000"
 
 
-def make_chapters(chapters: List[Chapter], timecodes: List[float], outfile: str) -> None:
-    chapters.sort(key=lambda c: c.frame)
-    rand = RandMan()
-
-    root = etree.Element("Chapters")
-    ed = etree.SubElement(root, "EditionEntry")
-    etree.SubElement(ed, "EditionUID").text = rand.get_rand()
+def chapters_xml(chapters: List[Chapter], timecodes: List[float]) -> List[_Element]:
+    atoms: List[_Element] = []
 
     for i, c in enumerate(chapters):
         start = timecode_to_timestamp(timecodes[c.frame])
-        end = timecode_to_timestamp(timecodes[chapters[i+1].frame]) if i < len(chapters) - 1 else None
-        atom = etree.SubElement(ed, "ChapterAtom")
+        end = timecode_to_timestamp(timecodes[c.end_frame]) if c.end_frame else \
+            timecode_to_timestamp(timecodes[chapters[i+1].frame]) if i < len(chapters) - 1 \
+            else None
+        atom = etree.Element("ChapterAtom")
         etree.SubElement(atom, "ChapterTimeStart").text = start
         if end is not None:
             etree.SubElement(atom, "ChapterTimeEnd").text = end
@@ -58,6 +67,51 @@ def make_chapters(chapters: List[Chapter], timecodes: List[float], outfile: str)
         etree.SubElement(disp, "ChapLanguageIETF").text = LANGMAP[c.lang]
         etree.SubElement(disp, "ChapterLanguage").text = c.lang
         etree.SubElement(atom, "ChapterUID").text = rand.get_rand()
+        atoms.append(atom)
+
+    return atoms
+
+
+def edition_xml(edition: Edition, timecodes: List[float]) -> _Element:
+    ed = etree.Element("EditionEntry")
+    etree.SubElement(ed, "EditionFlagOrdered").text = "1" if edition.ordered else "0"
+    etree.SubElement(ed, "EditionUID").text = rand.get_rand()
+    etree.SubElement(ed, "EditionFlagDefault").text = "1" if edition.default else "0"
+
+    for x in chapters_xml(edition.chapters, timecodes):
+        ed.append(x)
+
+    return ed
+
+
+def _to_edition(chapters: Optional[List[Chapter]] = None,
+                editions: Optional[List[Edition]] = None) -> List[Edition]:
+    if (chapters and editions) or (not chapters and not editions):
+        raise ValueError("Must supply one and only one of a list of chapters or list of editions!")
+
+    if chapters:
+        editions = [Edition(chapters=chapters, default=True, ordered=False)]
+
+    assert editions is not None
+    return editions
+
+
+def make_chapters(timecodes: List[float], outfile: str,
+                  chapters: Optional[List[Chapter]] = None,
+                  editions: Optional[List[Edition]] = None) -> None:
+    editions = _to_edition(chapters=chapters, editions=editions)
+    root = etree.Element("Chapters")
+    for e in editions:
+        root.append(edition_xml(e, timecodes))
 
     with open(outfile, "wb") as o:
         o.write(etree.tostring(root, encoding="utf-8", xml_declaration=True, pretty_print=True))
+
+
+def make_qpfile(qpfile: str,
+                chapters: Optional[List[Chapter]] = None,
+                editions: Optional[List[Edition]] = None) -> None:
+    editions = _to_edition(chapters=chapters, editions=editions)
+    frames = set(c.frame for e in editions for c in e.chapters)
+    with open(qpfile, "w", encoding="utf-8") as qp:
+        qp.writelines([f"{f} I\n" for f in sorted(list(frames))])

yt_common/yt_common/source.py

@@ -64,10 +64,14 @@ def glob_filename(pattern: str) -> str:
 class FileSource(ABC):
     def _open(self, path: str) -> vs.VideoNode:
         return depth(core.lsmas.LWLibavSource(path), 16) if path.lower().endswith(".m2ts") \
+            else depth(core.d2v.Source(path), 16) if path.lower().endswith(".d2v") \
             else depth(core.ffms2.Source(path), 16)
 
+    def audio_ref(self) -> Optional[vs.VideoNode]:
+        return None
+
     @abstractmethod
-    def get_audio(self) -> List[FileTrim]:
+    def audio_src(self) -> List[FileTrim]:
         pass
 
     @abstractmethod
@@ -77,16 +81,30 @@ class FileSource(ABC):
 
 class SimpleSource(FileSource):
     src: List[FileTrim]
+    sclip: Optional[vs.VideoNode]
+    aref: Optional[vs.VideoNode]
+    asrc: Optional[List[FileTrim]]
 
-    def __init__(self, src: Union[str, List[str], FileTrim, List[FileTrim]]) -> None:
+    def __init__(self, src: Union[str, List[str], FileTrim, List[FileTrim]],
+                 aref: Optional[vs.VideoNode] = None,
+                 asrc: Optional[Union[FileTrim, List[FileTrim]]] = None) -> None:
         srcl = src if isinstance(src, list) else [src]
         self.src = [FileTrim(s, (None, None)) if isinstance(s, str) else s for s in srcl]
+        self.sclip = None
+        self.aref = aref
+        self.asrc = asrc if isinstance(asrc, list) else [asrc] if asrc is not None else None
 
-    def get_audio(self) -> List[FileTrim]:
-        return self.src
+    def audio_ref(self) -> Optional[vs.VideoNode]:
+        return self.aref
+
+    def audio_src(self) -> List[FileTrim]:
+        return self.asrc if self.asrc else self.src
 
     def source(self) -> vs.VideoNode:
-        return core.std.Splice([s.apply_trim(self._open(s.path)) for s in self.src])
+        if self.sclip:
+            return self.sclip
+        self.sclip = core.std.Splice([s.apply_trim(self._open(s.path)) for s in self.src])
+        return self.sclip
 
 
 class DehardsubFileFinder(FileSource):
@@ -124,7 +142,7 @@ class FunimationSource(DehardsubFileFinder):
         self.ref_is_funi = False
         super().__init__(*args, **kwargs)
 
-    def get_audio(self) -> List[FileTrim]:
+    def audio_src(self) -> List[FileTrim]:
         if self.ref_is_funi:
             return [FileTrim(self.get_funi_filename(), (FUNI_INTRO, None))]