From d46b23a4f5fdceebf02cd61665eaeb1df62b4724 Mon Sep 17 00:00:00 2001 From: vikingowl Date: Mon, 11 Aug 2025 09:35:29 +0200 Subject: [PATCH] [refactor] extract and centralize output writing logic into `write_outputs` function in `output.rs` for improved code reuse and maintainability --- src/main.rs | 106 +++++++------------------------------------------- src/output.rs | 89 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 93 deletions(-) create mode 100644 src/output.rs diff --git a/src/main.rs b/src/main.rs index b585d68..35053dc 100644 --- a/src/main.rs +++ b/src/main.rs @@ -10,6 +10,9 @@ use clap::{Parser, Subcommand}; use clap_complete::Shell; use serde::{Deserialize, Serialize}; +mod output; +use output::{write_outputs, OutputFormats}; + use std::sync::mpsc::channel; // whisper-rs is used from the library crate use polyscribe::backend::{BackendKind, select_backend}; @@ -123,8 +126,8 @@ struct InputSegment { use polyscribe::{OutputEntry, date_prefix, models_dir_path, normalize_lang_code, render_srt}; #[derive(Debug, Serialize)] -struct OutputRoot { - items: Vec, +pub struct OutputRoot { + pub items: Vec, } fn sanitize_speaker_name(raw: &str) -> String { @@ -530,29 +533,8 @@ fn run() -> Result<()> { .unwrap_or("output"); let date = date_prefix(); let base_name = format!("{date}_{stem}"); - let json_path = out_dir.join(format!("{}.json", &base_name)); - let toml_path = out_dir.join(format!("{}.toml", &base_name)); - let srt_path = out_dir.join(format!("{}.srt", &base_name)); - - let mut json_file = File::create(&json_path).with_context(|| { - format!("Failed to create output file: {}", json_path.display()) - })?; - serde_json::to_writer_pretty(&mut json_file, &out)?; - writeln!(&mut json_file)?; - - let toml_str = toml::to_string_pretty(&out)?; - let mut toml_file = File::create(&toml_path).with_context(|| { - format!("Failed to create output file: {}", toml_path.display()) - })?; - toml_file.write_all(toml_str.as_bytes())?; - if !toml_str.ends_with('\n') { - writeln!(&mut toml_file)?; - } - - let srt_str = render_srt(&out.items); - let mut srt_file = File::create(&srt_path) - .with_context(|| format!("Failed to create output file: {}", srt_path.display()))?; - srt_file.write_all(srt_str.as_bytes())?; + let base_path = out_dir.join(&base_name); + write_outputs(&base_path, &out, &OutputFormats::all())?; // Extend merged with per-file entries merged_entries.extend(out.items.into_iter()); @@ -586,27 +568,8 @@ fn run() -> Result<()> { let date = date_prefix(); let merged_base = format!("{date}_merged"); - let m_json = out_dir.join(format!("{}.json", &merged_base)); - let m_toml = out_dir.join(format!("{}.toml", &merged_base)); - let m_srt = out_dir.join(format!("{}.srt", &merged_base)); - - let mut mj = File::create(&m_json) - .with_context(|| format!("Failed to create output file: {}", m_json.display()))?; - serde_json::to_writer_pretty(&mut mj, &merged_out)?; - writeln!(&mut mj)?; - - let m_toml_str = toml::to_string_pretty(&merged_out)?; - let mut mt = File::create(&m_toml) - .with_context(|| format!("Failed to create output file: {}", m_toml.display()))?; - mt.write_all(m_toml_str.as_bytes())?; - if !m_toml_str.ends_with('\n') { - writeln!(&mut mt)?; - } - - let m_srt_str = render_srt(&merged_out.items); - let mut ms = File::create(&m_srt) - .with_context(|| format!("Failed to create output file: {}", m_srt.display()))?; - ms.write_all(m_srt_str.as_bytes())?; + let base_path = out_dir.join(&merged_base); + write_outputs(&base_path, &merged_out, &OutputFormats::all())?; // Final concise summary table to stderr (below progress bars) if !args.quiet && !summary.is_empty() { @@ -775,29 +738,8 @@ fn run() -> Result<()> { let date = date_prefix(); let base_name = format!("{date}_{stem}"); let dir = parent_opt.unwrap_or(Path::new("")); - let json_path = dir.join(format!("{}.json", &base_name)); - let toml_path = dir.join(format!("{}.toml", &base_name)); - let srt_path = dir.join(format!("{}.srt", &base_name)); - - let mut json_file = File::create(&json_path).with_context(|| { - format!("Failed to create output file: {}", json_path.display()) - })?; - serde_json::to_writer_pretty(&mut json_file, &out)?; - writeln!(&mut json_file)?; - - let toml_str = toml::to_string_pretty(&out)?; - let mut toml_file = File::create(&toml_path).with_context(|| { - format!("Failed to create output file: {}", toml_path.display()) - })?; - toml_file.write_all(toml_str.as_bytes())?; - if !toml_str.ends_with('\n') { - writeln!(&mut toml_file)?; - } - - let srt_str = render_srt(&out.items); - let mut srt_file = File::create(&srt_path) - .with_context(|| format!("Failed to create output file: {}", srt_path.display()))?; - srt_file.write_all(srt_str.as_bytes())?; + let base_path = dir.join(&base_name); + write_outputs(&base_path, &out, &OutputFormats::all())?; } else { let stdout = io::stdout(); let mut handle = stdout.lock(); @@ -955,30 +897,8 @@ fn run() -> Result<()> { .unwrap_or("output"); let date = date_prefix(); let base_name = format!("{date}_{stem}"); - let json_path = dir.join(format!("{}.json", &base_name)); - let toml_path = dir.join(format!("{}.toml", &base_name)); - let srt_path = dir.join(format!("{}.srt", &base_name)); - - let mut json_file = File::create(&json_path).with_context(|| { - format!("Failed to create output file: {}", json_path.display()) - })?; - serde_json::to_writer_pretty(&mut json_file, &out)?; - writeln!(&mut json_file)?; - - let toml_str = toml::to_string_pretty(&out)?; - let mut toml_file = File::create(&toml_path).with_context(|| { - format!("Failed to create output file: {}", toml_path.display()) - })?; - toml_file.write_all(toml_str.as_bytes())?; - if !toml_str.ends_with('\n') { - writeln!(&mut toml_file)?; - } - - let srt_str = render_srt(&out.items); - let mut srt_file = File::create(&srt_path).with_context(|| { - format!("Failed to create output file: {}", srt_path.display()) - })?; - srt_file.write_all(srt_str.as_bytes())?; + let base_path = dir.join(&base_name); + write_outputs(&base_path, &out, &OutputFormats::all())?; } else { // stdout (only single input reaches here) let stdout = io::stdout(); diff --git a/src/output.rs b/src/output.rs new file mode 100644 index 0000000..3f4420f --- /dev/null +++ b/src/output.rs @@ -0,0 +1,89 @@ +use std::fs::File; +use std::io::Write; +use std::path::Path; + +use anyhow::Context; + +use crate::render_srt; +use crate::OutputRoot; + +/// Which formats to write. +pub struct OutputFormats { + pub json: bool, + pub toml: bool, + pub srt: bool, +} + +impl OutputFormats { + pub fn all() -> Self { + Self { json: true, toml: true, srt: true } + } +} + +/// Write outputs for the given base path (without extension). +/// This will create files named `base.json`, `base.toml`, and `base.srt` +/// according to the `formats` flags. JSON and TOML will always end with a trailing newline. +pub fn write_outputs(base: &Path, root: &OutputRoot, formats: &OutputFormats) -> anyhow::Result<()> { + if formats.json { + let json_path = base.with_extension("json"); + let mut json_file = File::create(&json_path).with_context(|| { + format!("Failed to create output file: {}", json_path.display()) + })?; + serde_json::to_writer_pretty(&mut json_file, root)?; + // ensure trailing newline + writeln!(&mut json_file)?; + } + + if formats.toml { + let toml_path = base.with_extension("toml"); + let toml_str = toml::to_string_pretty(root)?; + let mut toml_file = File::create(&toml_path).with_context(|| { + format!("Failed to create output file: {}", toml_path.display()) + })?; + toml_file.write_all(toml_str.as_bytes())?; + if !toml_str.ends_with('\n') { + writeln!(&mut toml_file)?; + } + } + + if formats.srt { + let srt_path = base.with_extension("srt"); + let srt_str = render_srt(&root.items); + let mut srt_file = File::create(&srt_path).with_context(|| { + format!("Failed to create output file: {}", srt_path.display()) + })?; + srt_file.write_all(srt_str.as_bytes())?; + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::OutputEntry; + + #[test] + fn write_outputs_creates_files_and_newlines() { + let dir = tempfile::tempdir().unwrap(); + let base = dir.path().join("test_base"); + let items = vec![OutputEntry { id: 0, speaker: "Alice".to_string(), start: 0.0, end: 1.23, text: "Hello".to_string() }]; + let root = OutputRoot { items }; + + write_outputs(&base, &root, &OutputFormats::all()).unwrap(); + + let json_path = base.with_extension("json"); + let toml_path = base.with_extension("toml"); + let srt_path = base.with_extension("srt"); + + assert!(json_path.exists(), "json file should exist"); + assert!(toml_path.exists(), "toml file should exist"); + assert!(srt_path.exists(), "srt file should exist"); + + let json = std::fs::read_to_string(&json_path).unwrap(); + let toml = std::fs::read_to_string(&toml_path).unwrap(); + + assert!(json.ends_with('\n'), "json should end with newline"); + assert!(toml.ends_with('\n'), "toml should end with newline"); + } +}