diff --git a/.vscode/settings.json b/.vscode/settings.json index 3f2ccb1..637231b 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,3 +1,3 @@ { - "rust-analyzer.cargo.features": ["rss"] + "rust-analyzer.cargo.features": ["rss", "indicatif", "audiotag"] } diff --git a/Cargo.toml b/Cargo.toml index ee401fd..220e5da 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -74,7 +74,10 @@ tracing-test = "0.2.5" # Included crates rustypipe = { path = ".", version = "0.2.0", default-features = false } -rustypipe-downloader = { path = "./downloader", version = "0.1.0", default-features = false } +rustypipe-downloader = { path = "./downloader", version = "0.1.0", default-features = false, features = [ + "indicatif", + "audiotag", +] } [features] default = ["default-tls"] diff --git a/README.md b/README.md index 432fa4c..15f6876 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,9 @@ # ![RustyPipe](https://code.thetadev.de/ThetaDev/rustypipe/raw/branch/main/notes/logo.svg) +[![Current crates.io version](https://img.shields.io/crates/v/smartcrop2.svg)](https://crates.io/crates/smartcrop2) +[![License](https://img.shields.io/badge/License-GPL--3-blue.svg?style=flat)](http://opensource.org/licenses/MIT) +[![CI status](https://code.thetadev.de/ThetaDev/rustypipe/actions/workflows/ci.yaml/badge.svg?style=flat&label=CI)](https://code.thetadev.de/ThetaDev/rustypipe/actions/?workflow=ci.yaml) + Rust client for the public YouTube / YouTube Music API (Innertube), inspired by [NewPipe](https://github.com/TeamNewPipe/NewPipeExtractor). diff --git a/cli/src/main.rs b/cli/src/main.rs index 484c51b..a4a428a 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -329,6 +329,8 @@ async fn download_video( .rustypipe(rp) .stream_filter(filter) .progress_bar(multi) + .audio_tag() + .crop_cover() .build(); let mut q = target.apply(dl.download_id(id)); if let Some(player_type) = player_type { @@ -361,6 +363,8 @@ async fn download_videos( .rustypipe(rp) .stream_filter(filter) .progress_bar(multi.clone()) + .audio_tag() + .crop_cover() .path_precheck() .build(); diff --git a/downloader/Cargo.toml b/downloader/Cargo.toml index c48fbe8..949a410 100644 --- a/downloader/Cargo.toml +++ b/downloader/Cargo.toml @@ -30,6 +30,8 @@ rustls-tls-native-roots = [ "rustypipe/rustls-tls-native-roots", ] +audiotag = ["dep:lofty", "dep:image", "dep:smartcrop2"] + [dependencies] rustypipe.workspace = true once_cell.workspace = true @@ -39,6 +41,10 @@ futures.workspace = true reqwest = { workspace = true, features = ["stream"] } rand.workspace = true tokio = { workspace = true, features = ["macros", "fs", "process"] } -indicatif.workspace = true +indicatif = { workspace = true, optional = true } filenamify.workspace = true tracing.workspace = true +time.workspace = true +lofty = { version = "0.21.0", optional = true } +image = { version = "0.25.0", optional = true } +smartcrop2 = { version = "0.2.0", optional = true } diff --git a/downloader/src/lib.rs b/downloader/src/lib.rs index 6e9cfdf..3b72541 100644 --- a/downloader/src/lib.rs +++ b/downloader/src/lib.rs @@ -8,6 +8,7 @@ use std::{ borrow::Cow, cmp::Ordering, ffi::OsString, + io::Cursor, ops::Range, path::{Path, PathBuf}, sync::Arc, @@ -15,7 +16,6 @@ use std::{ }; use futures::stream::{self, StreamExt}; -use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; use once_cell::sync::Lazy; use rand::Rng; use regex::Regex; @@ -23,17 +23,25 @@ use reqwest::{header, Client, StatusCode}; use rustypipe::{ client::{ClientType, RustyPipe}, model::{ + richtext::ToPlaintext, traits::{FileFormat, YtEntity}, - AudioCodec, VideoCodec, VideoPlayer, + AudioCodec, TrackItem, VideoCodec, VideoDetails, VideoPlayer, }, param::StreamFilter, }; +use time::{Date, OffsetDateTime}; use tokio::{ fs::{self, File}, io::AsyncWriteExt, process::Command, }; +#[cfg(feature = "indicatif")] +use indicatif::{MultiProgress, ProgressBar, ProgressStyle}; + +#[cfg(feature = "audiotag")] +use lofty::{config::WriteOptions, picture::Picture, prelude::*, tag::Tag}; + pub use util::DownloadError; type Result = core::result::Result; @@ -54,11 +62,16 @@ pub struct Downloader { pub struct DownloaderBuilder { rp: Option, ffmpeg: String, + #[cfg(feature = "indicatif")] multi: Option, filter: StreamFilter, video_format: DownloadVideoFormat, n_retries: u32, path_precheck: bool, + #[cfg(feature = "audiotag")] + audio_tag: bool, + #[cfg(feature = "audiotag")] + crop_cover: bool, } struct DownloaderInner { @@ -69,6 +82,7 @@ struct DownloaderInner { /// Path to the ffmpeg binary ffmpeg: String, /// Global progress + #[cfg(feature = "indicatif")] multi: Option, /// Default stream filter filter: StreamFilter, @@ -78,6 +92,12 @@ struct DownloaderInner { n_retries: u32, /// Check if destination path exists before player is fetched path_precheck: bool, + /// Apply metadata to audio files + #[cfg(feature = "audiotag")] + audio_tag: bool, + /// Crop YT thumbnails to ensure square album covers + #[cfg(feature = "audiotag")] + crop_cover: bool, } /// Download query @@ -89,6 +109,7 @@ pub struct DownloadQuery { /// Destination dest: DownloadDest, /// Progress bar + #[cfg(feature = "indicatif")] multi: Option, /// Stream filter filter: Option, @@ -200,11 +221,16 @@ impl Default for DownloaderBuilder { Self { rp: None, ffmpeg: "ffmpeg".to_owned(), + #[cfg(feature = "indicatif")] multi: None, filter: StreamFilter::new(), video_format: DownloadVideoFormat::Mp4, n_retries: 3, path_precheck: false, + #[cfg(feature = "audiotag")] + audio_tag: false, + #[cfg(feature = "audiotag")] + crop_cover: false, } } } @@ -235,6 +261,7 @@ impl DownloaderBuilder { /// Set the indicatif [`MultiProgress`] used to show download progress /// for all downloads + #[cfg(feature = "indicatif")] #[must_use] pub fn progress_bar(mut self, progress: MultiProgress) -> Self { self.multi = Some(progress); @@ -277,6 +304,22 @@ impl DownloaderBuilder { self } + /// Enable audio tagging + #[cfg(feature = "audiotag")] + #[must_use] + pub fn audio_tag(mut self) -> Self { + self.audio_tag = true; + self + } + + /// Crop YouTube thumbnails to get square album covers + #[cfg(feature = "audiotag")] + #[must_use] + pub fn crop_cover(mut self) -> Self { + self.crop_cover = true; + self + } + /// Create a new, configured [`Downloader`] instance pub fn build(self) -> Downloader { self.build_with_client( @@ -294,11 +337,16 @@ impl DownloaderBuilder { rp: self.rp.unwrap_or_default(), http: http_client, ffmpeg: self.ffmpeg, + #[cfg(feature = "indicatif")] multi: self.multi, filter: self.filter, video_format: self.video_format, n_retries: self.n_retries, path_precheck: self.path_precheck, + #[cfg(feature = "audiotag")] + audio_tag: self.audio_tag, + #[cfg(feature = "audiotag")] + crop_cover: self.crop_cover, }), } } @@ -328,6 +376,7 @@ impl Downloader { dl: self.clone(), video, dest: DownloadDest::Default, + #[cfg(feature = "indicatif")] multi: None, filter: None, video_format: None, @@ -414,6 +463,7 @@ impl DownloadQuery { } /// Use a [`MultiProgress`] progress bar for all downloads + #[cfg(feature = "indicatif")] pub fn progress_bar(mut self, progress: MultiProgress) -> Self { self.multi = Some(progress); self @@ -443,16 +493,26 @@ impl DownloadQuery { let mut last_err = None; // Progress bar - let multi = self.multi.clone().or_else(|| self.dl.i.multi.clone()); - let pb = multi.map(|m| { - let pb = ProgressBar::new(1); - pb.set_style(ProgressStyle::with_template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})").unwrap() - .progress_chars("#>-")); - m.add(pb) - }); + #[cfg(feature = "indicatif")] + let pb = { + let multi = self.multi.clone().or_else(|| self.dl.i.multi.clone()); + multi.map(|m| { + let pb = ProgressBar::new(1); + pb.set_style(ProgressStyle::with_template("{msg}\n{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {bytes}/{total_bytes} ({bytes_per_sec}, {eta})").unwrap() + .progress_chars("#>-")); + m.add(pb) + }) + }; for n in 0..=self.dl.i.n_retries { - let err = match self.download_attempt(&pb, n).await { + let err = match self + .download_attempt( + n, + #[cfg(feature = "indicatif")] + &pb, + ) + .await + { Ok(res) => return Ok(res), Err(DownloadError::Http(e)) => { if !e.is_timeout() && e.status() != Some(StatusCode::FORBIDDEN) { @@ -472,7 +532,11 @@ impl DownloadQuery { Err(last_err.unwrap()) } - async fn download_attempt(&self, pb: &Option, n: u32) -> Result { + async fn download_attempt( + &self, + #[allow(unused_variables)] n: u32, + #[cfg(feature = "indicatif")] pb: &Option, + ) -> Result { let filter = self.filter.as_ref().unwrap_or(&self.dl.i.filter); let video_format = self.video_format.unwrap_or(self.dl.i.video_format); @@ -495,11 +559,13 @@ impl DownloadQuery { } } + #[cfg(feature = "indicatif")] let attempt_suffix = if n > 0 { format!(" (retry #{n})") } else { String::new() }; + #[cfg(feature = "indicatif")] if let Some(pb) = pb { pb.set_message(format!( "Fetching player data for {}{}", @@ -550,7 +616,7 @@ impl DownloadQuery { .details .channel_name .clone() - .or(details.map(|d| d.channel.name)), + .or(details.as_ref().map(|d| d.channel.name.to_owned())), }; let output_path = self.dest.get_dest_path(&pv).with_extension(extension); @@ -580,11 +646,20 @@ impl DownloadQuery { }); } + #[cfg(feature = "indicatif")] if let Some(pb) = pb { pb.set_message(format!("Downloading {name}{attempt_suffix}")) } - download_streams(&downloads, &self.dl.i.http, &user_agent, pb.clone()).await?; + download_streams( + &downloads, + &self.dl.i.http, + &user_agent, + #[cfg(feature = "indicatif")] + pb.clone(), + ) + .await?; + #[cfg(feature = "indicatif")] if let Some(pb) = &pb { pb.set_message(format!("Converting {name}")); pb.set_style( @@ -595,6 +670,25 @@ impl DownloadQuery { } convert_streams(&downloads, &output_path, &self.dl.i.ffmpeg, &name).await?; + + // Tag audio file + #[cfg(feature = "audiotag")] + if self.dl.i.audio_tag && video.is_none() { + let (details, track) = match details { + Some(d) => (d, self.dl.i.rp.query().music_details(&self.video.id).await?), + None => { + let q = self.dl.i.rp.query(); + tokio::try_join!( + q.video_details(&self.video.id), + q.music_details(&self.video.id) + )? + } + }; + self.apply_audio_tags(&output_path, details, track.track) + .await?; + } + + #[cfg(feature = "indicatif")] if let Some(pb) = pb { pb.disable_steady_tick(); } @@ -606,8 +700,9 @@ impl DownloadQuery { .collect::>() .await .into_iter() - .collect::>()?; + .collect::>()?; + #[cfg(feature = "indicatif")] if let Some(pb) = pb { pb.finish_and_clear(); } @@ -616,6 +711,107 @@ impl DownloadQuery { player_data, }) } + + #[cfg(feature = "audiotag")] + async fn apply_audio_tags( + &self, + file: &Path, + details: VideoDetails, + track: TrackItem, + ) -> Result<()> { + use std::num::NonZeroU32; + + use image::codecs::jpeg::JpegEncoder; + + let mut tagged_file = lofty::read_from_path(file)?; + let tag = match tagged_file.primary_tag_mut() { + Some(primary_tag) => primary_tag, + None => { + if let Some(first_tag) = tagged_file.first_tag_mut() { + first_tag + } else { + let tag_type = tagged_file.primary_tag_type(); + tagged_file.insert_tag(Tag::new(tag_type)); + + tagged_file.primary_tag_mut().unwrap() + } + } + }; + + let description = details.description.to_plaintext(); + + tag.set_album( + track + .album + .map(|b| b.name) + .unwrap_or_else(|| track.name.clone()), + ); + tag.set_artist( + track + .artists + .into_iter() + .next() + .map(|a| a.name) + .unwrap_or(details.channel.name), + ); + tag.set_title(track.name); + if let Some(release_date) = extract_yt_release_date(&description, details.publish_date) { + if let Ok(date_str) = release_date.format(&YMD_FORMAT) { + tag.insert_text(ItemKey::RecordingDate, date_str); + } + } + tag.set_comment(description); + + let thumbnail = track.cover.into_iter().max_by_key(|c| c.height); + if let Some(thumbnail) = thumbnail { + let resp = self + .dl + .i + .http + .get(thumbnail.url) + .send() + .await? + .error_for_status()?; + let img_type = resp + .headers() + .get(header::CONTENT_TYPE) + .and_then(|fmt| fmt.to_str().ok()) + .and_then(image::ImageFormat::from_mime_type); + let img_bts = resp.bytes().await?; + + let mut lofty_img = if self.dl.i.crop_cover { + let mut img = if let Some(fmt) = img_type { + image::load_from_memory_with_format(&img_bts, fmt)? + } else { + image::load_from_memory(&img_bts)? + }; + + // Crop cover image if it is not square + if img.height() != img.width() && img.height() > 0 { + let crop = smartcrop::find_best_crop(&img, NonZeroU32::MIN, NonZeroU32::MIN) + .unwrap() + .crop; + img = img.crop_imm(crop.x, crop.y, crop.width, crop.height); + let mut enc_bts = Vec::new(); + img.write_with_encoder(JpegEncoder::new_with_quality(&mut enc_bts, 90))?; + let mut rd = Cursor::new(enc_bts); + Picture::from_reader(&mut rd)? + } else { + let mut rd = Cursor::new(img_bts); + Picture::from_reader(&mut rd)? + } + } else { + let mut rd = Cursor::new(img_bts); + Picture::from_reader(&mut rd)? + }; + + lofty_img.set_pic_type(lofty::picture::PictureType::CoverFront); + tag.set_picture(0, lofty_img); + } + + tag.save_to_path(file, WriteOptions::default())?; + Ok(()) + } } fn get_download_range(offset: u64, size: Option) -> Range { @@ -666,12 +862,12 @@ fn filenamify_lim(name: &str) -> String { } } -async fn download_single_file>( +async fn download_single_file( url: &str, - output: P, + output: &Path, http: &Client, user_agent: &str, - pb: Option, + #[cfg(feature = "indicatif")] pb: Option, ) -> Result<()> { // Check if file is already downloaded let output_path: PathBuf = output.into(); @@ -729,6 +925,7 @@ async fn download_single_file>( size = Some(original_size); offset = file_size; + #[cfg(feature = "indicatif")] if let Some(pb) = &pb { pb.inc_length(original_size); pb.inc(offset); @@ -759,10 +956,29 @@ async fn download_single_file>( .await?; if is_gvideo && size.is_some() { - download_chunks_by_param(http, &mut file, url, size.unwrap(), offset, user_agent, pb) - .await?; + download_chunks_by_param( + http, + &mut file, + url, + size.unwrap(), + offset, + user_agent, + #[cfg(feature = "indicatif")] + pb, + ) + .await?; } else { - download_chunks_by_header(http, &mut file, url, size, offset, user_agent, pb).await?; + download_chunks_by_header( + http, + &mut file, + url, + size, + offset, + user_agent, + #[cfg(feature = "indicatif")] + pb, + ) + .await?; } fs::rename(&output_path_tmp, &output_path).await?; @@ -779,7 +995,7 @@ async fn download_chunks_by_header( size: Option, offset: u64, user_agent: &str, - pb: Option, + #[cfg(feature = "indicatif")] pb: Option, ) -> Result<()> { let mut offset = offset; let mut size = size; @@ -826,6 +1042,7 @@ async fn download_chunks_by_header( offset = parsed_offset + 1; if size.is_none() { size = Some(parsed_size); + #[cfg(feature = "indicatif")] if let Some(pb) = &pb { pb.inc_length(parsed_size); } @@ -836,6 +1053,7 @@ async fn download_chunks_by_header( while let Some(item) = stream.next().await { // Retrieve chunk. let mut chunk = item?; + #[cfg(feature = "indicatif")] if let Some(pb) = &pb { pb.inc(chunk.len() as u64); } @@ -859,9 +1077,10 @@ async fn download_chunks_by_param( size: u64, offset: u64, user_agent: &str, - pb: Option, + #[cfg(feature = "indicatif")] pb: Option, ) -> Result<()> { let mut offset = offset; + #[cfg(feature = "indicatif")] if let Some(pb) = &pb { pb.inc_length(size); } @@ -891,6 +1110,7 @@ async fn download_chunks_by_param( while let Some(item) = stream.next().await { // Retrieve chunk. let mut chunk = item?; + #[cfg(feature = "indicatif")] if let Some(pb) = &pb { pb.inc(chunk.len() as u64); } @@ -918,12 +1138,21 @@ async fn download_streams( downloads: &Vec, http: &Client, user_agent: &str, - pb: Option, + #[cfg(feature = "indicatif")] pb: Option, ) -> Result<()> { let n = downloads.len(); stream::iter(downloads) - .map(|d| download_single_file(&d.url, d.file.clone(), http, user_agent, pb.clone())) + .map(|d| { + download_single_file( + &d.url, + &d.file, + http, + user_agent, + #[cfg(feature = "indicatif")] + pb.clone(), + ) + }) .buffer_unordered(n) .collect::>() .await @@ -933,9 +1162,9 @@ async fn download_streams( Ok(()) } -async fn convert_streams>( +async fn convert_streams( downloads: &[StreamDownload], - output: P, + output: &Path, ffmpeg: &str, title: &str, ) -> Result<()> { @@ -975,3 +1204,33 @@ async fn convert_streams>( } Ok(()) } + +const YMD_FORMAT: &[time::format_description::FormatItem] = + time::macros::format_description!("[year]-[month]-[day]"); + +fn extract_yt_release_date( + description: &str, + publish_date: Option, +) -> Option { + static RELEASE_DATE_REGEX: Lazy = + Lazy::new(|| Regex::new(r"Released on: (\d{4}-\d{2}-\d{2})").unwrap()); + + RELEASE_DATE_REGEX + .captures(description) + .and_then(|cap| { + let raw_date = &cap[1]; + Date::parse(raw_date, YMD_FORMAT).ok() + }) + .map(|release_date| { + if let Some(upload_date) = publish_date { + // Prefer the video upload date if it lies within 4 days of the release date + let upload_date = upload_date.date(); + let diff = (upload_date - release_date).abs(); + if diff < time::Duration::days(4) { + return upload_date; + } + } + release_date + }) + .or_else(|| publish_date.map(|d| d.date())) +} diff --git a/downloader/src/util.rs b/downloader/src/util.rs index d95b8f0..3934e2f 100644 --- a/downloader/src/util.rs +++ b/downloader/src/util.rs @@ -27,11 +27,29 @@ pub enum DownloadError { /// Download target already exists #[error("file {0} already exists")] Exists(PathBuf), + #[cfg(feature = "audiotag")] + /// Audio tagging error + #[error("Audio tag error: {0}")] + AudioTag(Cow<'static, str>), /// Other error #[error("error: {0}")] Other(Cow<'static, str>), } +#[cfg(feature = "audiotag")] +impl From for DownloadError { + fn from(value: lofty::error::LoftyError) -> Self { + Self::AudioTag(value.to_string().into()) + } +} + +#[cfg(feature = "audiotag")] +impl From for DownloadError { + fn from(value: image::ImageError) -> Self { + Self::AudioTag(value.to_string().into()) + } +} + /// Split an URL into its base string and parameter map /// /// Example: