ytdl-sub/ytdl_subscribe/downloaders/downloader.py
2022-04-22 20:34:44 +00:00

109 lines
3.7 KiB
Python

import abc
from abc import ABC
from contextlib import contextmanager
from pathlib import Path
from typing import Dict
from typing import Generic
from typing import List
from typing import Optional
from typing import Type
from typing import TypeVar
import yt_dlp as ytdl
from ytdl_subscribe.entries.entry import Entry
from ytdl_subscribe.validators.strict_dict_validator import StrictDictValidator
class DownloaderValidator(StrictDictValidator, ABC):
"""
Placeholder class to define downloader options
"""
DownloaderOptionsT = TypeVar("DownloaderOptionsT", bound=DownloaderValidator)
DownloaderEntryT = TypeVar("DownloaderEntryT", bound=Entry)
class Downloader(Generic[DownloaderOptionsT, DownloaderEntryT], ABC):
"""
Class that interacts with ytdl to perform the download of metadata and content,
and should translate that to list of Entry objects.
"""
downloader_options_type: Type[DownloaderOptionsT] = NotImplemented
downloader_entry_type: Type[DownloaderEntryT] = NotImplemented
@classmethod
def ytdl_option_overrides(cls) -> Dict:
"""Global overrides that even overwrite user input"""
return {"writethumbnail": True, "noplaylist": True}
@classmethod
def ytdl_option_defaults(cls) -> Dict:
"""Downloader defaults that can be overwritten from user input"""
return {}
@classmethod
def _configure_ytdl_options(
cls,
working_directory: str,
ytdl_options: Optional[Dict],
download_archive_file_name: Optional[str],
) -> Dict:
"""Configure the ytdl options for the downloader"""
if ytdl_options is None:
ytdl_options = {}
# Overwrite defaults with input
ytdl_options = dict(cls.ytdl_option_defaults(), **ytdl_options)
# Overwrite defaults + input with global options
ytdl_options = dict(ytdl_options, **cls.ytdl_option_overrides())
# Overwrite the output location with the specified working directory
ytdl_options["outtmpl"] = str(Path(working_directory) / "%(id)s.%(ext)s")
# If a download archive file name is provided, set it to that
ytdl_options["download_archive"] = str(Path(working_directory) / download_archive_file_name)
return ytdl_options
def __init__(
self,
working_directory: str,
download_options: DownloaderOptionsT,
ytdl_options: Optional[Dict] = None,
download_archive_file_name: Optional[str] = None,
):
self.working_directory = working_directory
self.download_options = download_options
self.ytdl_options = Downloader._configure_ytdl_options(
ytdl_options=ytdl_options,
working_directory=self.working_directory,
download_archive_file_name=download_archive_file_name,
)
@contextmanager
def ytdl_downloader(self, ytdl_options_overrides: Optional[Dict] = None) -> ytdl.YoutubeDL:
"""
Context manager to interact with yt_dlp.
"""
ytdl_options = self.ytdl_options
if ytdl_options_overrides is not None:
ytdl_options = dict(ytdl_options, **ytdl_options_overrides)
with ytdl.YoutubeDL(ytdl_options) as ytdl_downloader:
yield ytdl_downloader
def extract_info(self, ytdl_options_overrides: Optional[Dict] = None, **kwargs) -> Dict:
"""
Wrapper around yt_dlp.YoutubeDL.YoutubeDL.extract_info
All kwargs will passed to the extract_info function.
"""
with self.ytdl_downloader(ytdl_options_overrides) as ytdl_downloader:
return ytdl_downloader.extract_info(**kwargs)
@abc.abstractmethod
def download(self) -> List[DownloaderEntryT]:
"""The function to perform the download"""