improve video downloader, format sorting

This commit is contained in:
ThetaDev 2022-08-08 22:03:49 +02:00
parent 45bd435d34
commit 0c30794cb7
4 changed files with 159 additions and 58 deletions

View file

@ -11,5 +11,5 @@ tokio = {version = "1.20.0", features = ["rt-multi-thread"]}
indicatif = "0.17.0"
futures = "0.3.21"
anyhow = "1.0"
env_logger = "0.9.0"
clap = { version = "3.2.16", features = ["derive"] }
rusty-tube = {path = "../"}

View file

@ -1,11 +1,36 @@
use std::path::PathBuf;
use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use futures::stream::{self, StreamExt};
use indicatif::{MultiProgress, ProgressBar, ProgressStyle};
use reqwest::{Client, ClientBuilder};
use rusty_tube::client::{ClientType, RustyTube};
#[derive(Parser)]
#[clap(author, version, about, long_about = None)]
struct Cli {
#[clap(subcommand)]
command: Commands,
#[clap(short, value_parser, default_value = ".", global = true)]
output: PathBuf,
#[clap(long, value_parser, global = true)]
resolution: Option<u32>,
}
#[derive(Subcommand)]
enum Commands {
/// Download a playlist
Playlist {
/// Playlist ID
#[clap(value_parser)]
id: String,
},
}
async fn download_video(
video_id: String,
video_title: String,
output_dir: &str,
output_fname: Option<String>,
resolution: Option<u32>,
@ -19,37 +44,44 @@ async fn download_video(
pb.set_style(ProgressStyle::default_bar()
.template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})").unwrap()
.progress_chars("#>-"));
pb.set_message("Fetching player data");
pb.set_message(format!("Fetching player data for {}", video_title));
let player_data = rt
.get_player(video_id.as_str(), ClientType::Desktop)
let res = async {
let player_data = rt
.get_player(video_id.as_str(), ClientType::Desktop)
.await
.context(format!(
"Failed to fetch player data for video {}",
video_id
))?;
rusty_tube::download::download_video(
&player_data,
output_dir,
output_fname,
resolution,
ffmpeg,
http,
pb,
)
.await
.context(format!(
"Failed to fetch player data for video {}",
video_id
))?;
rusty_tube::download::download_video(
&player_data,
output_dir,
output_fname,
resolution,
ffmpeg,
http,
pb,
)
.await
.context(format!(
"Failed to download video '{}' [{}]",
player_data.info.title, video_id
))?;
"Failed to download video '{}' [{}]",
player_data.info.title, video_id
))
}
.await;
main.inc(1);
Ok(())
res
}
#[tokio::main]
async fn main() {
async fn download_playlist(
id: &str,
output_dir: &str,
output_fname: Option<String>,
resolution: Option<u32>,
) {
let http = ClientBuilder::new()
.user_agent("Mozilla/5.0 (Windows NT 10.0; Win64; rv:107.0) Gecko/20100101 Firefox/107.0")
.gzip(true)
@ -58,10 +90,7 @@ async fn main() {
.expect("unable to build the HTTP client");
let rt = RustyTube::new();
let playlist = rt
.get_playlist("PL4fGSI1pDJn4X-OicSCOy-dChXWdTgziQ", ClientType::Desktop)
.await
.unwrap();
let playlist = rt.get_playlist(id, ClientType::Desktop).await.unwrap();
// Indicatif setup
let multi = MultiProgress::new();
@ -70,7 +99,7 @@ async fn main() {
));
main.set_style(
ProgressStyle::default_bar()
.template("Downloading {pos:>}/{len} Videos [{wide_bar:.blue}]")
.template("Downloaded {pos:>}/{len} Videos [{wide_bar:.blue}]")
.unwrap()
.progress_chars("#>-"),
);
@ -80,9 +109,10 @@ async fn main() {
.map(|item| {
download_video(
item.video_id.to_owned(),
"tmp",
None,
None,
item.title.to_owned(),
output_dir,
output_fname.to_owned(),
resolution,
"ffmpeg",
&rt,
http.clone(),
@ -94,6 +124,52 @@ async fn main() {
.collect::<Vec<_>>()
.await
.into_iter()
.collect::<Result<(), _>>()
.unwrap();
.for_each(|res| match res {
Ok(_) => {}
Err(e) => {
println!("ERROR: {:?}", e);
}
});
}
#[tokio::main]
async fn main() {
let cli = Cli::parse();
// Cases: Existing folder, non-existing file with existing parent folder,
// Error cases: non-existing parent folder, existing file
let output_path = std::fs::canonicalize(cli.output).unwrap();
if output_path.is_file() {
println!("Output file already exists");
return;
}
let (output_dir, output_fname) = if output_path.is_dir() {
(output_path.to_string_lossy().to_string(), None)
} else {
let output_dir_parent = output_path.parent().unwrap();
if !output_dir_parent.is_dir() {
println!(
"Parent folder {} does not exist",
output_dir_parent.to_string_lossy()
);
return;
}
(
output_dir_parent.to_string_lossy().to_string(),
Some(
output_path
.file_name()
.unwrap()
.to_string_lossy()
.to_string(),
),
)
};
match cli.command {
Commands::Playlist { id } => {
download_playlist(&id, &output_dir, output_fname, cli.resolution).await
}
};
}

View file

@ -323,6 +323,18 @@ fn cmp_video_streams(a: &VideoStream, b: &VideoStream) -> Ordering {
}
}
fn cmp_audio_streams(a: &AudioStream, b: &AudioStream) -> Ordering {
fn cmp_bitrate(s: &AudioStream) -> u32 {
match s.codec {
// Opus is more efficient
AudioCodec::Opus => (s.average_bitrate as f32 * 1.3) as u32,
_ => s.average_bitrate,
}
}
cmp_bitrate(a).cmp(&cmp_bitrate(b))
}
fn map_player_data(response: response::Player, deobf: &Deobfuscator) -> Result<PlayerData> {
// Check playability status
match response.playability_status {
@ -423,7 +435,7 @@ fn map_player_data(response: response::Player, deobf: &Deobfuscator) -> Result<P
// Sort streams by quality
video_streams.sort_by(cmp_video_streams);
video_only_streams.sort_by(cmp_video_streams);
audio_streams.sort_by_key(|s| s.average_bitrate);
audio_streams.sort_by(cmp_audio_streams);
let subtitles = response.captions.map_or(vec![], |captions| {
captions
@ -629,7 +641,8 @@ mod tests {
&mut url_params,
&DEOBFUSCATOR,
&mut ["".to_owned(), "".to_owned()],
).unwrap();
)
.unwrap();
let url = Url::parse_with_params(url_base.as_str(), url_params.iter())
.unwrap()
.to_string();

View file

@ -4,7 +4,7 @@ use anyhow::{anyhow, bail, Result};
use fancy_regex::Regex;
use futures::stream::{self, StreamExt};
use indicatif::ProgressBar;
use log::debug;
use log::{debug, info};
use once_cell::sync::Lazy;
use rand::Rng;
use reqwest::{header, Client};
@ -189,6 +189,7 @@ pub async fn download_video(
// Download filepath
let download_dir = PathBuf::from(output_dir);
let title = player_data.info.title.to_owned();
let output_fname_set = output_fname.is_some();
let output_fname = output_fname
.unwrap_or_else(|| filenamify::filenamify(format!("{} [{}]", title, player_data.info.id)));
@ -199,7 +200,7 @@ pub async fn download_video(
.video_only_streams
.iter()
.rev()
.find(|s| s.height == r && !s.hdr)
.find(|s| s.height <= r && !s.hdr)
.clone(),
Err(anyhow!("no video stream matching res"))
)),
@ -207,10 +208,36 @@ pub async fn download_video(
};
let audio = some_or_bail!(
player_data.audio_streams.iter().rev().next(),
player_data
.audio_streams
.iter()
.rev()
.filter(|a| a.codec != AudioCodec::Unknown)
.next(),
Err(anyhow!("no audio stream"))
);
let format = match video {
Some(_) => "mp4",
None => match audio.codec {
AudioCodec::Unknown => panic!(),
AudioCodec::Mp4a => "m4a",
AudioCodec::Opus => "opus",
},
};
let output_path = download_dir.join(&output_fname).with_extension(format);
if output_path.exists() {
// If the downloaded video already exists, only error if the download path was
// chosen explicitly.
if output_fname_set {
bail!("File {} already exists", output_path.to_string_lossy());
} else {
info!("Downloaded video {} already exists", output_path.to_string_lossy());
return Ok(());
}
}
let mut downloads: Vec<StreamDownload> = Vec::new();
video.map(|v| {
@ -236,8 +263,7 @@ pub async fn download_video(
download_streams(&downloads, http, pb.clone()).await?;
pb.set_message(format!("Converting {}", title));
// let output_file = download_dir.join(format!("{}.mp4", title_fname));
convert_streams(&downloads, download_dir.join(output_fname), ffmpeg).await?;
convert_streams(&downloads, output_path, ffmpeg).await?;
// Delete original files
stream::iter(&downloads)
@ -286,21 +312,7 @@ async fn convert_streams<P: Into<PathBuf>>(
output: P,
ffmpeg: &str,
) -> Result<()> {
let format = if downloads.len() == 1
&& downloads[0].video_codec.is_none()
&& downloads[0].audio_codec.is_some()
{
match downloads[0].audio_codec.unwrap_or_default() {
AudioCodec::Unknown => bail!("unknown audio codec"),
AudioCodec::Mp4a => "m4a",
AudioCodec::Opus => "opus",
}
} else {
"mp4"
};
let mut output: PathBuf = output.into();
output.set_extension(format);
let output_path: PathBuf = output.into();
let mut args: Vec<OsString> = vec![];
let mut mapping_args: Vec<OsString> = vec![];
@ -318,7 +330,7 @@ async fn convert_streams<P: Into<PathBuf>>(
args.push("-c".into());
args.push("copy".into());
args.push(output.into());
args.push(output_path.into());
let res = Command::new(ffmpeg).args(args).output().await?;