implement separate download logic for youtube and tiktok, allowing for better quality decisions. DI via enums
This commit is contained in:
mykola2312
parent
6599410768
commit
b0afa21511
2 changed files with 115 additions and 83 deletions
196
src/dl.rs
196
src/dl.rs
|
|
@ -1,12 +1,12 @@
|
|||
use std::fmt;
|
||||
use std::fs;
|
||||
use tracing::{event, Level};
|
||||
|
||||
use crate::bot::sanitize::{extract_url, parse_url};
|
||||
use crate::dl::ffmpeg::FFMpeg;
|
||||
|
||||
use self::spawn::SpawnError;
|
||||
use self::tmpfile::{TmpFile, TmpFileError};
|
||||
use self::yt_dlp::{YtDlp, YtDlpError, YtDlpFormat, YtDlpInfo};
|
||||
use self::yt_dlp::{YtDlp, YtDlpError, YtDlpInfo};
|
||||
|
||||
pub mod ffmpeg;
|
||||
mod spawn;
|
||||
|
|
@ -15,6 +15,7 @@ pub mod yt_dlp;
|
|||
|
||||
pub enum DownloadError {
|
||||
Message(String),
|
||||
NotAnURL,
|
||||
NoFormatFound,
|
||||
MakePathError,
|
||||
}
|
||||
|
|
@ -44,6 +45,7 @@ impl fmt::Display for DownloadError {
|
|||
use DownloadError as DE;
|
||||
match &self {
|
||||
DE::Message(msg) => write!(f, "{}", msg),
|
||||
DE::NotAnURL => write!(f, "no url or malformed url were provided"),
|
||||
DE::NoFormatFound => write!(
|
||||
f,
|
||||
"no best format found. you may want to specify one yourself"
|
||||
|
|
@ -53,102 +55,134 @@ impl fmt::Display for DownloadError {
|
|||
}
|
||||
}
|
||||
|
||||
fn make_download_path(
|
||||
info: &YtDlpInfo,
|
||||
suffix: Option<&str>,
|
||||
format: &YtDlpFormat,
|
||||
) -> Result<String, DownloadError> {
|
||||
std::env::temp_dir()
|
||||
.join(format!(
|
||||
"{}_{}.{}",
|
||||
info.id,
|
||||
suffix.unwrap_or(""),
|
||||
format.ext
|
||||
))
|
||||
.into_os_string()
|
||||
.into_string()
|
||||
.map_err(|e| DownloadError::MakePathError)
|
||||
enum Downloader {
|
||||
Default,
|
||||
YouTube,
|
||||
TikTok,
|
||||
}
|
||||
|
||||
fn file_exists(path: &str) -> bool {
|
||||
match fs::metadata(path) {
|
||||
Ok(_) => true,
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
const DEFAULT_DOWNLOADER: (&'static str, Downloader) = ("", Downloader::Default);
|
||||
const DOWNLOADERS: [(&'static str, Downloader); 4] = [
|
||||
("www.youtube.com", Downloader::YouTube),
|
||||
("youtu.be", Downloader::YouTube),
|
||||
("www.tiktok.com", Downloader::TikTok),
|
||||
("vm.tiktok.com", Downloader::TikTok),
|
||||
];
|
||||
|
||||
pub fn delete_if_exists(path: &str) {
|
||||
if file_exists(path) {
|
||||
if let Err(e) = fs::remove_file(path) {
|
||||
event!(Level::ERROR, "{}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn download_fallback(url: &str, info: YtDlpInfo) -> Result<TmpFile, DownloadError> {
|
||||
let av = match info.best_av_format() {
|
||||
Some(av) => av,
|
||||
None => {
|
||||
event!(
|
||||
Level::WARN,
|
||||
"no best format found for {}, reverting to default",
|
||||
url
|
||||
);
|
||||
match info.default_format() {
|
||||
Some(format) => format,
|
||||
None => {
|
||||
event!(Level::ERROR, "no formats found for {}", url);
|
||||
return Err(DownloadError::NoFormatFound);
|
||||
impl Downloader {
|
||||
async fn default_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
|
||||
let av = match info.best_av_format() {
|
||||
Some(av) => av,
|
||||
None => {
|
||||
event!(
|
||||
Level::WARN,
|
||||
"no best format found for {}, reverting to default",
|
||||
url
|
||||
);
|
||||
match info.default_format() {
|
||||
Some(format) => format,
|
||||
None => {
|
||||
event!(Level::ERROR, "no formats found for {}", url);
|
||||
return Err(DownloadError::NoFormatFound);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
};
|
||||
|
||||
Ok(YtDlp::download(url, &info, &av).await?)
|
||||
}
|
||||
Ok(YtDlp::download(url, &info, &av).await?)
|
||||
}
|
||||
|
||||
pub async fn download(url: &str) -> Result<TmpFile, DownloadError> {
|
||||
event!(Level::INFO, "url {}", url);
|
||||
async fn youtube_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
|
||||
let vf = match info.best_video_format() {
|
||||
Some(vf) => vf,
|
||||
None => return Err(DownloadError::NoFormatFound),
|
||||
};
|
||||
let af = match info.best_audio_format() {
|
||||
Some(af) => af,
|
||||
None => return Err(DownloadError::NoFormatFound),
|
||||
};
|
||||
|
||||
let info = YtDlp::load_info(url).await?;
|
||||
let vf = match info.best_video_format() {
|
||||
Some(vf) => vf,
|
||||
None => return download_fallback(url, info).await,
|
||||
};
|
||||
let af = match info.best_audio_format() {
|
||||
Some(af) => af,
|
||||
None => return download_fallback(url, info).await,
|
||||
};
|
||||
let video = YtDlp::download(url, &info, &vf).await?;
|
||||
let audio = YtDlp::download(url, &info, &af).await?;
|
||||
|
||||
let video = YtDlp::download(url, &info, &vf).await?;
|
||||
let audio = YtDlp::download(url, &info, &af).await?;
|
||||
let abr = if let Some(abr) = af.abr {
|
||||
FFMpeg::round_mp3_bitrate(abr)
|
||||
} else {
|
||||
event!(
|
||||
Level::ERROR,
|
||||
"somehow url {} audio format {} doesnt have abr",
|
||||
url,
|
||||
af.format_id
|
||||
);
|
||||
|
||||
let abr = if let Some(abr) = af.abr {
|
||||
FFMpeg::round_mp3_bitrate(abr)
|
||||
} else {
|
||||
192
|
||||
};
|
||||
|
||||
let output = TmpFile::new(format!("{}.{}", &info.id, &vf.ext).as_str())?;
|
||||
event!(
|
||||
Level::ERROR,
|
||||
"somehow url {} audio format {} doesnt have abr",
|
||||
Level::INFO,
|
||||
"for {} we joining video {} and audio {}",
|
||||
url,
|
||||
vf.format_id,
|
||||
af.format_id
|
||||
);
|
||||
|
||||
192
|
||||
};
|
||||
let res = FFMpeg::join_video_audio(&video.path, &audio.path, abr, &output.path).await;
|
||||
|
||||
let output = TmpFile::new(format!("{}.{}", &info.id, &vf.ext).as_str())?;
|
||||
event!(
|
||||
Level::INFO,
|
||||
"for {} we joining video {} and audio {}",
|
||||
url,
|
||||
vf.format_id,
|
||||
af.format_id
|
||||
);
|
||||
match res {
|
||||
Ok(()) => Ok(output),
|
||||
Err(e) => Err(DownloadError::Message(e.to_string())),
|
||||
}
|
||||
}
|
||||
|
||||
let res = FFMpeg::join_video_audio(&video.path, &audio.path, abr, &output.path).await;
|
||||
async fn tiktok_download(url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
|
||||
let original = info.formats
|
||||
.iter()
|
||||
.find(|f| f.format_id == "0")
|
||||
.ok_or(DownloadError::NoFormatFound)?;
|
||||
|
||||
match res {
|
||||
Ok(()) => Ok(output),
|
||||
Err(e) => Err(DownloadError::Message(e.to_string())),
|
||||
Ok(YtDlp::download(url, info, original).await?)
|
||||
}
|
||||
|
||||
pub async fn download(&self, url: &str, info: &YtDlpInfo) -> Result<TmpFile, DownloadError> {
|
||||
match self {
|
||||
Downloader::Default => Self::default_download(url, info).await,
|
||||
Downloader::YouTube => Self::youtube_download(url, info).await,
|
||||
Downloader::TikTok => Self::tiktok_download(url, info).await
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl fmt::Display for Downloader {
|
||||
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
match self {
|
||||
Downloader::Default => write!(f, "Default"),
|
||||
Downloader::YouTube => write!(f, "YouTube"),
|
||||
Downloader::TikTok => write!(f, "TikTok")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn download(url: &str) -> Result<TmpFile, DownloadError> {
|
||||
let url = parse_url(extract_url(url).ok_or(DownloadError::NotAnURL)?)
|
||||
.ok_or(DownloadError::NotAnURL)?;
|
||||
let host_url = url.host_str().ok_or(DownloadError::NotAnURL)?;
|
||||
|
||||
let downloader = &DOWNLOADERS
|
||||
.iter()
|
||||
.find(|f| f.0 == host_url)
|
||||
.unwrap_or(&DEFAULT_DOWNLOADER).1;
|
||||
event!(Level::INFO, "using {} downloader for {}", downloader, url);
|
||||
|
||||
let info = YtDlp::load_info(url.as_str()).await?;
|
||||
let output = match downloader.download(url.as_str(), &info).await {
|
||||
Ok(output) => output,
|
||||
Err(e) => {
|
||||
event!(Level::ERROR, "downloader {} failed: {}. falling back to default downloader", downloader, e);
|
||||
|
||||
DEFAULT_DOWNLOADER.1.download(url.as_str(), &info).await?
|
||||
}
|
||||
};
|
||||
|
||||
Ok(output)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -241,8 +241,6 @@ impl fmt::Display for YtDlpError {
|
|||
}
|
||||
|
||||
pub struct YtDlp {}
|
||||
|
||||
// BUG: REAL ARGUMENT INJECTION! FIX ASAP
|
||||
impl YtDlp {
|
||||
pub async fn load_info(url: &str) -> Result<YtDlpInfo, YtDlpError> {
|
||||
let output = spawn("python", &["-m", "yt_dlp", url, "-j"]).await?;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue