improve video downloader, format sorting

cli/Cargo.toml

@@ -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 = "../"}

cli/src/main.rs

@@ -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
+    let res = async {
-        .get_player(video_id.as_str(), ClientType::Desktop)
+        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 {}",
+            "Failed to download video '{}' [{}]",
-            video_id
+            player_data.info.title, video_id
-        ))?;
+        ))
-
+    }
-    rusty_tube::download::download_video(
+    .await;
-        &player_data,
-        output_dir,
-        output_fname,
-        resolution,
-        ffmpeg,
-        http,
-        pb,
-    )
-    .await
-    .context(format!(
-        "Failed to download video '{}' [{}]",
-        player_data.info.title, video_id
-    ))?;
 
     main.inc(1);
-    Ok(())
+    res
 }
 
-#[tokio::main]
+async fn download_playlist(
-async fn main() {
+    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
+    let playlist = rt.get_playlist(id, ClientType::Desktop).await.unwrap();
-        .get_playlist("PL4fGSI1pDJn4X-OicSCOy-dChXWdTgziQ", 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",
+                item.title.to_owned(),
-                None,
+                output_dir,
-                None,
+                output_fname.to_owned(),
+                resolution,
                 "ffmpeg",
                 &rt,
                 http.clone(),
@@ -94,6 +124,52 @@ async fn main() {
         .collect::<Vec<_>>()
         .await
         .into_iter()
-        .collect::<Result<(), _>>()
+        .for_each(|res| match res {
-        .unwrap();
+            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
+        }
+    };
 }

src/client/player.rs

@@ -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();

src/download.rs

@@ -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, output_path, ffmpeg).await?;
-    convert_streams(&downloads, download_dir.join(output_fname), 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
+    let output_path: PathBuf = output.into();
-        && 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 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?;