subtitle_translator_cli/
processor.rs

1use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
2
3use crate::{
4    get_all_files, pb_init, read_file_trim_bom, sort_and_extract_translations, SrtFile,
5    SubtitleFile,
6};
7use core::time;
8use std::error::Error;
9use std::fs::{self, File};
10use std::io::Write;
11use std::sync::{Arc, Mutex};
12use std::thread;
13
14pub fn process_file(
15    file_path: String,
16    input_language: String,
17    output_language: String,
18) -> Result<(), Box<dyn Error>> {
19    // when file_path == "*.srt", should get all files in the current directory
20    let mut files = Vec::new();
21    if file_path.contains("*") {
22        files = get_all_files(file_path.as_str()).unwrap();
23    } else {
24        files.push(std::path::PathBuf::from(file_path));
25    }
26    for file in files {
27        process_single_file(
28            file.to_str().unwrap(),
29            input_language.clone(),
30            output_language.clone(),
31        )
32        .unwrap();
33    }
34    Ok(())
35}
36
37fn process_single_file(
38    file_path: &str,
39    input_language: String,
40    output_language: String,
41) -> Result<(), Box<dyn Error>> {
42    let contents = fs::read_to_string(&file_path).expect("Something went wrong reading the file");
43    let contents = read_file_trim_bom(&contents);
44    let file_name = file_path.split('.').next().unwrap_or("");
45    let file_extension = file_path.split('.').last().unwrap_or("");
46
47    let file_struct = match file_extension {
48        "srt" => Box::new(SrtFile {}),
49        _ => return Err("Unsupported file type".into()),
50    };
51
52    let split_contents = file_struct.split_contents(&contents).unwrap();
53
54    let translated_combined_text = run_translation_tasks(
55        split_contents,
56        input_language.clone(),
57        output_language.clone(),
58    );
59
60    let merged_contents = file_struct
61        .merge_contents(&contents, translated_combined_text)
62        .unwrap();
63
64    let mut file = File::create(format!(
65        "{}_{}.{}",
66        file_name,
67        output_language.clone(),
68        file_extension
69    ))?;
70
71    file.write_all(merged_contents.as_bytes())?;
72    Ok(())
73}
74
75fn run_translation_tasks(
76    contents: Vec<String>,
77    input_language: String,
78    output_language: String,
79) -> Vec<String> {
80    let translated_combined_text = Arc::new(Mutex::new(Vec::with_capacity(contents.len())));
81    spawn_translation_threads(
82        contents,
83        input_language,
84        output_language,
85        translated_combined_text.clone(),
86    );
87
88    let mut translated_combined_text = Arc::try_unwrap(translated_combined_text)
89        .expect("Arc unwrap failed")
90        .into_inner()
91        .expect("Failed to acquire lock");
92
93    sort_and_extract_translations(&mut translated_combined_text)
94}
95
96fn spawn_translation_threads(
97    contents: Vec<String>,
98    input_language: String,
99    output_language: String,
100    translated_combined_text: Arc<Mutex<Vec<(usize, String)>>>,
101) {
102    let progress = Arc::new(Mutex::new(0));
103    let total = contents.len(); // 总任务数
104    let pb = pb_init(total as u64);
105
106    contents
107        .into_par_iter()
108        .enumerate()
109        .for_each(|(index, content)| {
110            let translated_text = crate::translate(
111                content.clone(),
112                input_language.clone(),
113                output_language.clone(),
114            );
115
116            {
117                let mut translated_combined_text = translated_combined_text
118                    .lock()
119                    .expect("Failed to acquire lock");
120                translated_combined_text.push((index, translated_text));
121            }
122
123            {
124                let mut progress = progress.lock().expect("Failed to acquire lock");
125                *progress += 1;
126                pb.set_position(*progress);
127                if *progress % 5 == 0 {
128                    thread::sleep(time::Duration::from_millis(200));
129                }
130            }
131        });
132    pb.finish_with_message("finished");
133}
subtitle_translator_cli/processor.rs

subtitle_translator_cli/
processor.rs
