From d4e91277e3e28d37bee3120814400d23c77e138c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Florian=20Mei=C3=9Fner?= <936176+t-nil@users.noreply.github.com> Date: Tue, 16 Jan 2024 11:31:23 +0100 Subject: [PATCH] add choice on output file existing --- Cargo.lock | 94 ++++++++++++++++++++++++- Cargo.toml | 20 +++++- src/command/args.rs | 34 ++++++++- src/command/encode.rs | 158 ++++++++++++++++++++++++++++++++++++++++-- src/process.rs | 5 +- 5 files changed, 300 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1c5839d..6f95d48 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -18,14 +18,18 @@ dependencies = [ "humantime", "indicatif", "infer", + "mktemp", "serde", "serde_json", "shell-escape", "sled", + "strum", + "strum_macros", "time", "tokio", "tokio-process-stream", "tokio-stream", + "unix-named-pipe", ] [[package]] @@ -317,6 +321,17 @@ version = "0.3.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a357d28ed41a50f9c765dbfe56cbc04a64e53e5fc58ba79fbc34c10ef3df831f" +[[package]] +name = "errno" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f639046355ee4f37944e44f60642c6f3a7efa3cf6b78c78a0d989a8ce6c396a1" +dependencies = [ + "errno-dragonfly", + "libc", + "winapi", +] + [[package]] name = "errno" version = "0.3.8" @@ -327,6 +342,16 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "errno-dragonfly" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa68f1b12764fab894d2755d2518754e71b4fd80ecfb822714a1206c2aab39bf" +dependencies = [ + "cc", + "libc", +] + [[package]] name = "fastrand" version = "2.0.1" @@ -591,6 +616,15 @@ dependencies = [ "windows-sys 0.48.0", ] +[[package]] +name = "mktemp" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69fed8fbcd01affec44ac226784c6476a6006d98d13e33bc0ca7977aaf046bd8" +dependencies = [ + "uuid", +] + [[package]] name = "num_cpus" version = "1.16.0" @@ -731,12 +765,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "322394588aaf33c24007e8bb3238ee3e4c5c09c084ab32bc73890b99ff326bca" dependencies = [ "bitflags 2.4.1", - "errno", + "errno 0.3.8", "libc", "linux-raw-sys", "windows-sys 0.52.0", ] +[[package]] +name = "rustversion" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ffc183a10b4478d04cbbbfc96d0873219d962dd5accaff2ffbd4ceb7df837f4" + [[package]] name = "ryu" version = "1.0.16" @@ -826,12 +866,44 @@ version = "1.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2593d31f82ead8df961d8bd23a64c2ccf2eb5dd34b0a34bfb4dd54011c72009e" +[[package]] +name = "socket2" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5fac59a5cb5dd637972e5fca70daf0523c9067fcdc4842f053dae04a18f8e9" +dependencies = [ + "libc", + "windows-sys 0.48.0", +] + [[package]] name = "strsim" version = "0.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "73473c0e59e6d5812c5dfe2a064a6444949f089e20eec9a2e5506596494e4623" +[[package]] +name = "strum" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290d54ea6f91c969195bdbcd7442c8c2a2ba87da8bf60a7ee86a235d4bc1e125" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.25.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23dc1fa9ac9c169a78ba62f0b841814b7abae11bdd047b9c58f893439e309ea0" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn", +] + [[package]] name = "syn" version = "2.0.48" @@ -914,6 +986,7 @@ dependencies = [ "num_cpus", "pin-project-lite", "signal-hook-registry", + "socket2", "tokio-macros", "windows-sys 0.48.0", ] @@ -980,12 +1053,31 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e51733f11c9c4f72aa0c160008246859e340b00807569a0da0e7a1079b27ba85" +[[package]] +name = "unix-named-pipe" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ad653da8f36ac5825ba06642b5a3cce14a4e52c6a5fab4a8928d53f4426dae2" +dependencies = [ + "errno 0.2.8", + "libc", +] + [[package]] name = "utf8parse" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "711b9620af191e0cdc7468a8d14e709c3dcdb115b36f838e601583af800a370a" +[[package]] +name = "uuid" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79daa5ed5740825c40b389c5e50312b9c86df53fccd33f281df655642b43869d" +dependencies = [ + "getrandom", +] + [[package]] name = "wasi" version = "0.11.0+wasi-snapshot-preview1" diff --git a/Cargo.toml b/Cargo.toml index 42543f4..5953f7c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,11 +26,29 @@ serde = { version = "1.0.185", features = ["derive"] } serde_json = "1.0.105" shell-escape = "0.1.5" sled = "0.34.7" +strum = { version = "0.25.0", features = ["derive"] } +strum_macros = "0.25.3" time = { version = "0.3", features = ["parsing", "macros"] } -tokio = { version = "1.15", features = ["rt", "macros", "process", "fs", "signal"] } +tokio = { version = "1.15", features = [ + "rt", + "macros", + "process", + "fs", + "signal", +] } tokio-process-stream = "0.4" tokio-stream = "0.1" +[dev-dependencies] +mktemp = "0.5.1" + +[target.'cfg(unix)'.dependencies] +unix-named-pipe = "0.2" + +[target.'cfg(windows)'.dependencies.tokio] +version = "1.15" +features = ["net"] + [profile.release] lto = true opt-level = "s" diff --git a/src/command/args.rs b/src/command/args.rs index 91eab7d..310e731 100644 --- a/src/command/args.rs +++ b/src/command/args.rs @@ -3,10 +3,11 @@ mod encode; mod vmaf; pub use encode::*; +use strum::EnumCount; pub use vmaf::*; use crate::{command::encode::default_output_ext, ffprobe::Ffprobe}; -use clap::Parser; +use clap::{Parser, ValueEnum}; use std::{ path::{Path, PathBuf}, sync::Arc, @@ -22,6 +23,10 @@ pub struct EncodeToOutput { #[arg(short, long)] pub output: Option, + /// What to do if the output file already exists. + #[arg(long, default_value("rename"))] + pub on_duplicate: Option, + /// Set the output ffmpeg audio codec. /// By default 'copy' is used. Otherwise, if re-encoding is necessary, 'libopus' is default. /// @@ -95,3 +100,30 @@ impl Sample { self.extension = output.extension().and_then(|e| e.to_str().map(Into::into)); } } + +#[derive(Default, Clone, Copy, PartialEq, Eq, Debug, Hash, EnumCount)] +pub enum OnDuplicate { + Overwrite, + Skip, + #[default] + Rename, + Ask, +} + +impl ValueEnum for OnDuplicate { + fn value_variants<'a>() -> &'a [Self] { + let variants = &[Self::Overwrite, Self::Skip, Self::Rename, Self::Ask]; + // safety check in case one forgets to add a new variant + debug_assert!(OnDuplicate::COUNT == variants.len()); + variants + } + + fn to_possible_value(&self) -> Option { + match self { + OnDuplicate::Overwrite => Some(clap::builder::PossibleValue::new("overwrite")), + OnDuplicate::Skip => Some(clap::builder::PossibleValue::new("skip")), + OnDuplicate::Rename => Some(clap::builder::PossibleValue::new("rename").help("Tries to rename the output file by appending and incrementing a number.\nE.g. `vid.av1.mkv` -> `vid.av1_1.mkv` -> `vid.av1_2.mkv` etc.")), + OnDuplicate::Ask => Some(clap::builder::PossibleValue::new("ask")), + } + } +} diff --git a/src/command/encode.rs b/src/command/encode.rs index b792dbc..db92ea7 100644 --- a/src/command/encode.rs +++ b/src/command/encode.rs @@ -1,6 +1,6 @@ use crate::{ command::{ - args::{self, Encoder}, + args::{self, Encoder, OnDuplicate}, SmallDuration, PROGRESS_CHARS, }, console_ext::style, @@ -9,10 +9,12 @@ use crate::{ process::FfmpegOut, temporary::{self, TempKind}, }; +use anyhow::{anyhow, bail}; use clap::Parser; use console::style; use indicatif::{HumanBytes, ProgressBar, ProgressStyle}; use std::{ + io::Write, path::{Path, PathBuf}, sync::Arc, time::Duration, @@ -54,6 +56,7 @@ pub async fn run( encode: args::EncodeToOutput { output, + on_duplicate, audio_codec, downmix_to_stereo, video_only, @@ -62,17 +65,69 @@ pub async fn run( probe: Arc, bar: &ProgressBar, ) -> anyhow::Result<()> { - let defaulting_output = output.is_none(); // let probe = ffprobe::probe(&args.input); let output = output.unwrap_or_else(|| default_output_name(&args.input, &args.encoder, probe.is_image)); + let on_duplicate = on_duplicate.unwrap_or(OnDuplicate::default()); + + let output = match on_duplicate { + OnDuplicate::Overwrite => output, + OnDuplicate::Rename => rename_if_exists(output)?, + OnDuplicate::Skip => { + if output.exists() { + bar.println(format!( + "{} {}", + style("Skipping").dim(), + style(output.display()).dim(), + )); + bail!("Output file already exists"); + } else { + output + } + } + OnDuplicate::Ask => { + if output.exists() { + /*bar. + ));*/ + bar.suspend(|| loop { + let mut input = String::new(); + print!( + "{} {}. {}", + style("Output file already exists:"), + style(output.display()).dim(), + style("Overwrite, rename, skip, or quit? [o/r/s/q]").italic() + ); + std::io::stdout().flush()?; + std::io::stdin().read_line(&mut input)?; + match input.trim() { + "o" | "overwrite" => { + break Ok(output); + } + "r" | "rename" => { + break rename_if_exists(output); + } + "s" | "skip" => { + bail!("Output file already exists"); + } + "q" | "quit" => { + bail!("User quit"); + } + _ => { + eprintln!("Invalid input"); + } + } + })? + } else { + output + } + } + }; + // output is temporary until encoding has completed successfully temporary::add(&output, TempKind::NotKeepable); - if defaulting_output { - let out = shell_escape::escape(output.display().to_string().into()); - bar.println(style!("Encoding {out}").dim().to_string()); - } + let out = shell_escape::escape(output.display().to_string().into()); + bar.println(style!("Encoding {out}").dim().to_string()); bar.set_message("encoding, "); let mut enc_args = args.to_encoder_args(crf, &probe)?; @@ -145,6 +200,97 @@ pub async fn run( Ok(()) } +fn rename_if_exists(mut output: PathBuf) -> anyhow::Result { + while output.exists() { + // get basename without extension, or full name in case that fails + let name = output + .file_stem() + .or_else(|| output.file_name()) + .ok_or(anyhow!("Could not parse file name from {:?}", output))? + .to_string_lossy(); + + // if the last part of the file stem after an underscore is a valid positive integer, + // increment it by one. Otherwise add an "_1" suffix. + let mut parts: Vec = name.split('_').map(&str::to_owned).collect(); + if let Some(last_part) = parts.pop() { + if let Ok(number) = last_part.parse::() { + parts.push((number + 1).to_string()); + } else { + parts.push(last_part); + parts.push("1".to_owned()); + } + } else { + // this shouldn't happen since the name would have to be equal to "" (the empty string) + bail!("Output name vector {:?} should't be empty", parts); + } + + let name = parts.join("_") + + &output + .extension() + .map(|e| ".".to_owned() + &e.to_string_lossy()) + .unwrap_or("".to_owned()); + output = output.with_file_name(name); + } + Ok(output) +} + +#[test] +fn test_rename_if_exists() -> anyhow::Result<()> { + use anyhow::ensure; + use std::ops::Range; + use std::panic::catch_unwind; + + let temp_dir = mktemp::Temp::new_dir()?; + let mut temp_file = PathBuf::new(); + temp_file.push(&temp_dir); + + touch(&temp_file, "test.mkv")?; + _do(&temp_file, "test.mkv", "test_{}.mkv", 1..13)?; + + touch(&temp_dir, "test_1_2_3")?; + _do(&temp_file, "test_1_2_1", "test_1_2_{}", 1..3)?; + assert!(catch_unwind(|| { _do(&temp_file, "test_1_2_1", "test_1_2_{}", 3..4) }).is_err()); + _do(&temp_file, "test_1_2_1", "test_1_2_{}", 4..7)?; + + _do(&temp_file, "test_1_2", "test_1_{}", 2..6)?; + + touch(&temp_file, ".hidden_27")?; + _do(&temp_file, ".hidden", ".hidden", 0..1)?; + _do(&temp_file, ".hidden", ".hidden_{}", 1..27)?; + assert!(catch_unwind(|| { _do(&temp_file, ".hidden", ".hidden_{}", 27..28) }).is_err()); + _do(&temp_file, ".hidden", ".hidden_{}", 28..31)?; + + fn touch(f: &Path, name: &str) -> anyhow::Result<()> { + let mut p = PathBuf::new(); + p.push(f); + p.push(name); + std::process::Command::new("touch").arg(p).output()?; + Ok(()) + } + + fn _do(temp_dir: &Path, name: &str, pattern: &str, range: Range) -> anyhow::Result<()> { + ensure!(range.len() > 0, "range must be non-empty"); + + let mut temp_file = PathBuf::new(); + temp_file.push(temp_dir); + temp_file.push(name); + + for i in range { + let result = rename_if_exists(temp_file.clone())?; + let actual = result + .file_name() + .map(|s| s.to_string_lossy()) + .unwrap_or("".into()); + let expected = pattern.replace("{}", i.to_string().as_str()); + assert_eq!(actual, expected); + + std::process::Command::new("touch").arg(&result).output()?; + } + Ok(()) + } + Ok(()) +} + /// * vid.mp4 -> "mp4" /// * vid.??? -> "mkv" /// * image.??? -> "avif" diff --git a/src/process.rs b/src/process.rs index f47b965..f58a945 100644 --- a/src/process.rs +++ b/src/process.rs @@ -66,6 +66,7 @@ pub enum FfmpegOut { impl FfmpegOut { pub fn try_parse(line: &str) -> Option { if line.starts_with("frame=") { + eprintln!("{}", line); let frame: u64 = parse_label_substr("frame=", line)?.parse().ok()?; let fps: f32 = parse_label_substr("fps=", line)?.parse().ok()?; let (h, m, s, ns) = time::Time::parse( @@ -74,11 +75,11 @@ impl FfmpegOut { ) .ok()? .as_hms_nano(); - return Some(Self::Progress { + return Some(dbg!(Self::Progress { frame, fps, time: Duration::new(h as u64 * 60 * 60 + m as u64 * 60 + s as u64, ns), - }); + })); } if line.starts_with("video:") && line.contains("muxing overhead") { let video = parse_label_size("video:", line)?;