diff --git a/src/args.rs b/src/args.rs index 797f3f2..007c4b7 100644 --- a/src/args.rs +++ b/src/args.rs @@ -675,7 +675,11 @@ pub fn load_kirikiri_chat_json( let mut outt = arg.output_type.unwrap_or(OutputScriptType::M3t); if !matches!( outt, - OutputScriptType::M3t | OutputScriptType::M3ta | OutputScriptType::M3tTxt + OutputScriptType::M3t + | OutputScriptType::M3ta + | OutputScriptType::M3tTxt + | OutputScriptType::Po + | OutputScriptType::Pot ) { outt = OutputScriptType::M3t; } @@ -689,13 +693,21 @@ pub fn load_kirikiri_chat_json( &f, true, )?; - let m3t = crate::output_scripts::m3t::M3tParser::new( - &data, - arg.llm_trans_mark.as_ref().map(|s| s.as_str()), - ) - .parse_as_map()?; + let m3t = if outt.is_m3t() { + crate::output_scripts::m3t::M3tParser::new( + &data, + arg.llm_trans_mark.as_ref().map(|s| s.as_str()), + ) + .parse_as_map()? + } else { + crate::output_scripts::po::PoParser::new( + &data, + arg.llm_trans_mark.as_ref().map(|s| s.as_str()), + ) + .parse_as_map()? + }; for (k, v) in m3t { - map.insert(k, v); + map.insert(k.replace("\\[", "["), v.replace("\\[", "[")); } } return Ok(Some(std::sync::Arc::new(map))); diff --git a/src/output_scripts/po.rs b/src/output_scripts/po.rs index 504a9d6..c6041eb 100644 --- a/src/output_scripts/po.rs +++ b/src/output_scripts/po.rs @@ -589,12 +589,65 @@ impl<'a> PoParser<'a> { r } + pub fn parse_as_map(&mut self) -> Result> { + let mut map = HashMap::new(); + let mut llm = None; + for (i, entry) in self.parse_entries()?.into_iter().enumerate() { + if entry.msgid.is_empty() && i == 0 { + // This is the header entry, skip it + continue; + } + for comment in &entry.comments { + if let Comment::Translator(s) = comment { + let s = s.trim(); + if s.starts_with("NAME:") { + // name = Some(s[5..].trim().to_string()); + } else if s.starts_with("LLM:") { + llm = Some(s[4..].trim().replace("\\n", "\n")); + } + } + } + let message = match entry.msgstr { + MsgStr::Single(s) => { + let s = s.trim(); + if s.is_empty() { + llm.take() + .map(|mut llm| { + if let Some(mark) = self.llm_mark { + llm.push_str(mark); + } + llm + }) + .unwrap_or_else(|| { + String::from(if entry.msgid.is_empty() { "" } else { "" }) + }) + } else { + let mut tmp = s.to_string(); + if let Some(llm) = llm.take() { + if tmp == llm { + if let Some(mark) = self.llm_mark { + tmp.push_str(mark); + } + } + } + tmp + } + } + MsgStr::Plural(_) => { + return Err(anyhow!("Plural msgstr not supported in this context")); + } + }; + map.insert(entry.msgid, message); + } + Ok(map) + } + pub fn parse(&mut self) -> Result> { let mut messages = Vec::new(); let mut llm = None; let mut name = None; - for entry in self.parse_entries()? { - if entry.msgid.is_empty() { + for (i, entry) in self.parse_entries()?.into_iter().enumerate() { + if entry.msgid.is_empty() && i == 0 { // This is the header entry, skip it continue; } diff --git a/src/types.rs b/src/types.rs index 9bcb01f..cf35343 100644 --- a/src/types.rs +++ b/src/types.rs @@ -122,6 +122,14 @@ impl OutputScriptType { pub fn is_custom(&self) -> bool { matches!(self, OutputScriptType::Custom) } + + /// Returns true if the script type is M3t/M3ta/M3tTxt. + pub fn is_m3t(&self) -> bool { + matches!( + self, + OutputScriptType::M3t | OutputScriptType::M3ta | OutputScriptType::M3tTxt + ) + } } impl AsRef for OutputScriptType {