5
0
2021-05-02 23:20:27 -04:00

105 lines
3.1 KiB
Python

import vapoursynth as vs
import vsutil
from lvsfunc.types import Range
import lvsfunc as lvf
import glob
import os
from abc import ABC, abstractmethod
from typing import Any, List, Tuple
from .config import Config
from .logging import log
core = vs.core
SUBSPLS_FILENAME: str = "[SubsPlease] {title_long} - {epnum:02d} ({resolution}p) [$GLOB].mkv"
ER_FILENAME: str = "[Erai-raws] {title_long} - {epnum:02d} [v0][{resolution}p]$GLOB.mkv"
FUNI_INTRO: int = 289
AMAZON_FILENAME: str = "{title_long} - {epnum:02d} (Amazon Prime CBR {resolution}p).mkv"
def waka_replace(src: vs.VideoNode, wakas: List[vs.VideoNode], ranges: List[List[Range]]
) -> Tuple[vs.VideoNode, List[vs.VideoNode]]:
if len(wakas) == 0:
return src, wakas
new_wakas = []
for waka, r in zip(wakas, ranges):
tmp = src
src = lvf.misc.replace_ranges(src, waka, r)
new_wakas.append(lvf.misc.replace_ranges(waka, tmp, r))
return src, new_wakas
def glob_filename(pattern: str) -> str:
res = glob.glob(glob.escape(pattern).replace("$GLOB", "*"))
if len(res) == 0:
raise FileNotFoundError(f"File matching \"{pattern}\" not found!")
return res[0]
class DehardsubFileFinder(ABC):
config: Config
def __init__(self, config: Config) -> None:
self.config = config
@abstractmethod
def get_waka_filenames(self) -> List[str]:
pass
@abstractmethod
def get_ref(self) -> vs.VideoNode:
pass
def source(self) -> Tuple[List[vs.VideoNode], vs.VideoNode]:
wakas: List[vs.VideoNode] = []
for f in [self.config.format_filename(f) for f in self.get_waka_filenames()]:
if not os.path.isfile(f):
log.warn("Missing a waka!")
continue
wakas.append(vsutil.depth(core.ffms2.Source(f), 16))
ref = vsutil.depth(self.get_ref(), 16)
return wakas, ref
class FunimationSource(DehardsubFileFinder):
ref_is_funi: bool
def __init__(self, *args: Any, **kwargs: Any) -> None:
self.ref_is_funi = False
super().__init__(*args, **kwargs)
def get_amazon(self) -> vs.VideoNode:
if not os.path.isfile(self.config.format_filename(AMAZON_FILENAME)):
log.warn("Amazon not found, falling back to Funimation")
raise FileNotFoundError()
log.success("Found Amazon video")
return core.ffms2.Source(self.config.format_filename(AMAZON_FILENAME))
def get_funi_filename(self) -> str:
try:
return glob_filename(self.config.format_filename(ER_FILENAME))
except FileNotFoundError:
pass
try:
return glob_filename(self.config.format_filename(SUBSPLS_FILENAME))
except FileNotFoundError:
log.error("Could not find funimation video")
raise
def get_funi(self) -> vs.VideoNode:
return core.ffms2.Source(self.get_funi_filename())[FUNI_INTRO:]
def get_ref(self) -> vs.VideoNode:
try:
return self.get_amazon()
except FileNotFoundError:
self.ref_is_funi = True
return self.get_funi()