diff --git a/bdfr/site_downloaders/download_factory.py b/bdfr/site_downloaders/download_factory.py index 49dba5f..2c7ef21 100644 --- a/bdfr/site_downloaders/download_factory.py +++ b/bdfr/site_downloaders/download_factory.py @@ -18,6 +18,7 @@ from bdfr.site_downloaders.redgifs import Redgifs from bdfr.site_downloaders.self_post import SelfPost from bdfr.site_downloaders.vidble import Vidble from bdfr.site_downloaders.youtube import Youtube +from bdfr.site_downloaders.vreddit import vreddit class DownloadFactory: @@ -47,6 +48,8 @@ class DownloadFactory: return Youtube elif re.match(r'i\.redd\.it.*', sanitised_url): return Direct + elif re.match(r'v\.redd\.it.*', sanitised_url): + return vreddit elif re.match(r'pornhub\.com.*', sanitised_url): return PornHub elif re.match(r'vidble\.com', sanitised_url): diff --git a/bdfr/site_downloaders/vreddit.py b/bdfr/site_downloaders/vreddit.py new file mode 100644 index 0000000..945ee93 --- /dev/null +++ b/bdfr/site_downloaders/vreddit.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 + +import logging +import tempfile +from pathlib import Path +from typing import Callable, Optional + +import yt_dlp +from praw.models import Submission + +from bdfr.exceptions import NotADownloadableLinkError, SiteDownloaderError +from bdfr.resource import Resource +from bdfr.site_authenticator import SiteAuthenticator +from bdfr.site_downloaders.base_downloader import BaseDownloader + +logger = logging.getLogger(__name__) + + +class vreddit(BaseDownloader): + def __init__(self, post: Submission): + super().__init__(post) + + def find_resources(self, authenticator: Optional[SiteAuthenticator] = None) -> list[Resource]: + ytdl_options = { + 'playlistend': 1, + 'nooverwrites': True, + } + download_function = self._download_video(ytdl_options) + extension = self.get_video_attributes(self.post.url)['ext'] + res = Resource(self.post, self.post.url, download_function, extension) + return [res] + + def _download_video(self, ytdl_options: dict) -> Callable: + yt_logger = logging.getLogger('youtube-dl') + yt_logger.setLevel(logging.CRITICAL) + ytdl_options['quiet'] = True + ytdl_options['logger'] = yt_logger + + def download(_: dict) -> bytes: + with tempfile.TemporaryDirectory() as temp_dir: + download_path = Path(temp_dir).resolve() + ytdl_options['outtmpl'] = str(download_path) + '/' + 'test.%(ext)s' + try: + with yt_dlp.YoutubeDL(ytdl_options) as ydl: + ydl.download([self.post.url]) + except yt_dlp.DownloadError as e: + raise SiteDownloaderError(f'Vreddit download failed: {e}') + + downloaded_files = list(download_path.iterdir()) + if len(downloaded_files) > 0: + downloaded_file = downloaded_files[0] + else: + raise NotADownloadableLinkError(f"No media exists in the URL {self.post.url}") + with open(downloaded_file, 'rb') as file: + content = file.read() + return content + return download + + @staticmethod + def get_video_attributes(url: str) -> dict: + yt_logger = logging.getLogger('youtube-dl') + yt_logger.setLevel(logging.CRITICAL) + with yt_dlp.YoutubeDL({'logger': yt_logger, }) as ydl: + try: + result = ydl.extract_info(url, download=False) + except Exception as e: + logger.exception(e) + raise NotADownloadableLinkError(f'Video info extraction failed for {url}') + if 'ext' in result: + return result + else: + try: + result = (result["entries"][0]) + return result + except Exception as e: + logger.exception(e) + raise NotADownloadableLinkError(f'Video info extraction failed for {url}')