tinymist_l10n/
lib.rs

1//! Tinymist's localization library.
2
3use core::panic;
4use std::{
5    borrow::Cow,
6    collections::HashSet,
7    path::Path,
8    sync::{OnceLock, RwLock},
9};
10
11use rayon::{
12    iter::{IntoParallelRefMutIterator, ParallelIterator},
13    str::ParallelString,
14};
15use rustc_hash::FxHashMap;
16
17/// A map of translations.
18pub type TranslationMap = FxHashMap<String, String>;
19/// A set of translation maps.
20pub type TranslationMapSet = FxHashMap<String, TranslationMap>;
21
22static ALL_TRANSLATIONS: OnceLock<TranslationMapSet> = OnceLock::new();
23static LOCALE_TRANSLATIONS: RwLock<Option<&'static TranslationMap>> = RwLock::new(Option::None);
24
25/// Sets the current translations. It can only be called once.
26pub fn set_translations(translations: TranslationMapSet) {
27    let new_translations = ALL_TRANSLATIONS.set(translations);
28
29    if let Err(new_translations) = new_translations {
30        eprintln!("cannot set translations: len = {}", new_translations.len());
31    }
32}
33
34/// Sets the current locale.
35pub fn set_locale(locale: &str) -> Option<()> {
36    let translations = ALL_TRANSLATIONS.get()?;
37    let lower_locale = locale.to_lowercase();
38    let locale = lower_locale.as_str();
39    let translations = translations.get(locale).or_else(|| {
40        // Tries s to find a language that starts with the locale and follow a hyphen.
41        translations
42            .iter()
43            .find(|(k, _)| locale.starts_with(*k) && locale.chars().nth(k.len()) == Some('-'))
44            .map(|(_, v)| v)
45    })?;
46
47    *LOCALE_TRANSLATIONS.write().unwrap() = Some(translations);
48
49    Some(())
50}
51
52/// Loads a TOML string into a map of translations.
53pub fn load_translations(input: &str) -> anyhow::Result<TranslationMapSet> {
54    let mut translations = deserialize(input, false)?;
55    translations.par_iter_mut().for_each(|(_, v)| {
56        v.par_iter_mut().for_each(|(_, v)| {
57            if !v.starts_with('"') {
58                return;
59            }
60
61            *v = serde_json::from_str::<String>(v)
62                .unwrap_or_else(|e| panic!("cannot parse translation message: {e}, message: {v}"));
63        });
64    });
65
66    Ok(translations)
67}
68
69/// Updates disk translations with new key-value pairs.
70pub fn update_disk_translations(
71    mut key_values: Vec<(String, String)>,
72    output: &Path,
73) -> anyhow::Result<()> {
74    key_values.sort_by(|(key_x, _), (key_y, _)| key_x.cmp(key_y));
75
76    // Reads and parses existing translations
77    let mut translations = match std::fs::read_to_string(output) {
78        Ok(existing_translations) => deserialize(&existing_translations, true)?,
79        Err(e) if e.kind() == std::io::ErrorKind::NotFound => TranslationMapSet::default(),
80        Err(e) => Err(e)?,
81    };
82
83    // Removes unused translations
84    update_translations(key_values, &mut translations);
85
86    // Writes translations
87    let result = serialize_translations(translations);
88    std::fs::write(output, result)?;
89    Ok(())
90}
91
92/// Updates a map of translations with new key-value pairs.
93pub fn update_translations(
94    key_values: Vec<(String, String)>,
95    translations: &mut TranslationMapSet,
96) {
97    let used = key_values.iter().map(|e| &e.0).collect::<HashSet<_>>();
98    translations.retain(|k, _| used.contains(k));
99
100    // Updates translations
101    let en = "en".to_owned();
102    for (key, value) in key_values {
103        translations
104            .entry(key)
105            .or_default()
106            .insert(en.clone(), value);
107    }
108}
109
110/// Writes a map of translations to a TOML string.
111pub fn serialize_translations(translations: TranslationMapSet) -> String {
112    let mut result = String::new();
113
114    result.push_str("\n# The translations are partially generated by copilot\n");
115
116    let mut translations = translations.into_iter().collect::<Vec<_>>();
117    translations.sort_by(|a, b| a.0.cmp(&b.0));
118
119    for (key, mut data) in translations {
120        result.push_str(&format!("\n[{key}]\n"));
121
122        let en = data.remove("en").expect("en translation is missing");
123        result.push_str(&format!("en = {en}\n"));
124
125        // sort by lang
126        let mut data = data.into_iter().collect::<Vec<_>>();
127        data.sort_by(|a, b| a.0.cmp(&b.0));
128
129        for (lang, value) in data {
130            result.push_str(&format!("{lang} = {value}\n"));
131        }
132    }
133
134    result
135}
136
137/// Tries to translate a string to the current language.
138#[macro_export]
139macro_rules! t {
140    ($key:expr, $message:expr) => {
141        $crate::t_without_args($key, $message)
142    };
143    ($key:expr, $message:expr $(, $arg_key:ident = $arg_value:expr)+ $(,)?) => {
144        $crate::t_with_args($key, $message, &[$((stringify!($arg_key), $arg_value)),*])
145    };
146}
147
148/// Returns an error with a translated message.
149#[macro_export]
150macro_rules! bail {
151    ($key:expr, $message:expr $(, $arg_key:ident = $args:expr)* $(,)?) => {{
152        let msg = $crate::t!($key, $message $(, $arg_key = $args)*);
153        return Err(tinymist_std::error::prelude::_msg(concat!(file!(), ":", line!(), ":", column!()), msg.into()));
154    }};
155}
156
157/// Tries to get a translation for a key.
158fn find_message(key: &'static str, message: &'static str) -> &'static str {
159    let Some(translations) = LOCALE_TRANSLATIONS.read().unwrap().as_ref().copied() else {
160        return message;
161    };
162
163    translations.get(key).map(String::as_str).unwrap_or(message)
164}
165
166/// Tries to translate a string to the current language.
167pub fn t_without_args(key: &'static str, message: &'static str) -> Cow<'static, str> {
168    Cow::Borrowed(find_message(key, message))
169}
170
171/// An argument for a translation.
172pub enum Arg<'a> {
173    /// A string argument.
174    Str(Cow<'a, str>),
175    /// An integer argument.
176    Int(i64),
177    /// A float argument.
178    Float(f64),
179}
180
181impl<'a> From<&'a String> for Arg<'a> {
182    fn from(s: &'a String) -> Self {
183        Arg::Str(Cow::Borrowed(s.as_str()))
184    }
185}
186
187impl<'a> From<&'a str> for Arg<'a> {
188    fn from(s: &'a str) -> Self {
189        Arg::Str(Cow::Borrowed(s))
190    }
191}
192
193/// Converts an object to an argument of debug message.
194pub trait DebugL10n {
195    /// Returns a debug string for the current language.
196    fn debug_l10n(&self) -> Arg<'_>;
197}
198
199impl<T: std::fmt::Debug> DebugL10n for T {
200    fn debug_l10n(&self) -> Arg<'static> {
201        Arg::Str(Cow::Owned(format!("{self:?}")))
202    }
203}
204
205/// Tries to translate a string to the current language.
206pub fn t_with_args(
207    key: &'static str,
208    message: &'static str,
209    args: &[(&'static str, Arg)],
210) -> Cow<'static, str> {
211    let message = find_message(key, message);
212    let mut result = String::new();
213
214    // for c in message.chars() {
215    //     if c == '{' {
216    //         let Some(bracket_index) = message[arg_index..].find('}') else {
217    //             result.push(c);
218    //             continue;
219    //         };
220
221    //         let arg_index_str = &message[arg_index + 1..arg_index +
222    // bracket_index];
223
224    //         match arg {
225    //             Arg::Str(s) => result.push_str(s.as_ref()),
226    //             Arg::Int(i) => result.push_str(&i.to_string()),
227    //             Arg::Float(f) => result.push_str(&f.to_string()),
228    //         }
229
230    //         arg_index += arg_index_str.len() + 2;
231    //     } else {
232    //         result.push(c);
233    //     }
234    // }
235
236    let message_iter = &mut message.chars();
237    while let Some(c) = message_iter.next() {
238        if c == '{' {
239            let arg_index_str = message_iter.take_while(|c| *c != '}').collect::<String>();
240            message_iter.next();
241
242            let Some(arg) = args
243                .iter()
244                .find(|(k, _)| k == &arg_index_str)
245                .map(|(_, v)| v)
246            else {
247                result.push(c);
248                result.push_str(&arg_index_str);
249                continue;
250            };
251
252            match arg {
253                Arg::Str(s) => result.push_str(s.as_ref()),
254                Arg::Int(i) => result.push_str(&i.to_string()),
255                Arg::Float(f) => result.push_str(&f.to_string()),
256            }
257        } else {
258            result.push(c);
259        }
260    }
261
262    Cow::Owned(result)
263}
264
265/// Deserializes a TOML string into a map of translations.
266pub fn deserialize(input: &str, key_first: bool) -> anyhow::Result<TranslationMapSet> {
267    let lines = input
268        .par_split('\n')
269        .map(|line| line.trim())
270        .filter(|line| !line.starts_with('#') && !line.is_empty())
271        .collect::<Vec<_>>();
272
273    let mut translations = FxHashMap::default();
274    let mut key = String::new();
275
276    for line in lines {
277        if line.starts_with('[') {
278            key = line[1..line.len() - 1].to_string();
279        } else {
280            let equal_index = line.find('=').map_or_else(
281                || {
282                    Err(anyhow::anyhow!(
283                        "cannot find equal sign in translation line: {line}"
284                    ))
285                },
286                Ok,
287            )?;
288            let lang = line[..equal_index].trim().to_string();
289            let value = line[equal_index + 1..].trim().to_string();
290
291            if key_first {
292                translations
293                    .entry(key.clone())
294                    .or_insert_with(FxHashMap::default)
295                    .insert(lang, value);
296            } else {
297                translations
298                    .entry(lang)
299                    .or_insert_with(FxHashMap::default)
300                    .insert(key.clone(), value);
301            }
302        }
303    }
304
305    Ok(translations)
306}