unicode_width_utils/unicode_width_utils.rs
1use std::borrow::Cow;
2use std::sync::LazyLock;
3use std::sync::atomic::{AtomicBool, Ordering};
4use unicode_width::{UnicodeWidthChar, UnicodeWidthStr};
5
6static IS_CJK: LazyLock<AtomicBool> = LazyLock::new(|| {
7 let is_cjk = match std::env::var("UNICODE_WIDTH") {
8 Ok(value) => value.eq_ignore_ascii_case("cjk"),
9 _ => false,
10 };
11 AtomicBool::new(is_cjk)
12});
13
14/// A configuration helper to measure character and string widths.
15///
16/// It determines the width of Unicode characters and strings,
17/// optionally treating East Asian Ambiguous characters
18/// (such as certain Greek, Cyrillic, and CJK characters)
19/// as having a width of 2 (CJK mode).
20///
21/// The default CJK mode is initialized at startup
22/// based on the `UNICODE_WIDTH` environment variable,
23/// but can also be dynamically modified using [`UnicodeWidth::set_default_cjk`].
24#[derive(Clone, Debug)]
25pub struct UnicodeWidth {
26 is_cjk: bool,
27 should_expand_tab: bool,
28 tab_size: u8,
29}
30
31impl Default for UnicodeWidth {
32 /// Create a `UnicodeWidth` instance using the default CJK mode.
33 ///
34 /// The default CJK mode is determined by the global setting,
35 /// which defaults to the value of the `UNICODE_WIDTH` environment variable
36 /// (value `"cjk"` enabling CJK mode).
37 ///
38 /// See also [`set_default_cjk`].
39 ///
40 /// [`set_default_cjk`]: UnicodeWidth::set_default_cjk
41 fn default() -> Self {
42 Self {
43 is_cjk: IS_CJK.load(Ordering::Relaxed),
44 should_expand_tab: false,
45 tab_size: 0,
46 }
47 }
48}
49
50impl UnicodeWidth {
51 /// Create a `UnicodeWidth` instance using the default CJK mode.
52 ///
53 /// The default CJK mode is determined by the global setting,
54 /// which defaults to the value of the `UNICODE_WIDTH` environment variable
55 /// (value `"cjk"` enabling CJK mode).
56 ///
57 /// See also [`set_default_cjk`].
58 ///
59 /// [`set_default_cjk`]: UnicodeWidth::set_default_cjk
60 pub fn new() -> Self {
61 Self::default()
62 }
63
64 /// Create a new `UnicodeWidth` instance with a specific CJK flag.
65 ///
66 /// If `is_cjk` is true,
67 /// East Asian Ambiguous characters will be treated as 2 columns wide.
68 /// If false, they will be treated as 1 column wide.
69 ///
70 /// # Examples
71 /// ```
72 /// use unicode_width_utils::UnicodeWidth;
73 ///
74 /// let non_cjk = UnicodeWidth::with_cjk(false);
75 /// assert_eq!(non_cjk.char('█'), Some(1));
76 ///
77 /// let cjk = UnicodeWidth::with_cjk(true);
78 /// assert_eq!(cjk.char('█'), Some(2));
79 /// ```
80 pub fn with_cjk(is_cjk: bool) -> Self {
81 Self {
82 is_cjk,
83 ..Default::default()
84 }
85 }
86
87 /// Set the default CJK configuration dynamically.
88 ///
89 /// Future instances
90 /// created using [`UnicodeWidth::new`] or [`UnicodeWidth::default`]
91 /// will inherit this default value
92 /// unless explicitly overridden with [`UnicodeWidth::with_cjk`].
93 ///
94 /// # Examples
95 /// ```
96 /// use unicode_width_utils::UnicodeWidth;
97 ///
98 /// // Set default CJK mode to true.
99 /// UnicodeWidth::set_default_cjk(true);
100 /// let uw = UnicodeWidth::new();
101 /// assert_eq!(uw.char('█'), Some(2));
102 ///
103 /// // Set default CJK mode to false.
104 /// UnicodeWidth::set_default_cjk(false);
105 /// let uw2 = UnicodeWidth::new();
106 /// assert_eq!(uw2.char('█'), Some(1));
107 /// ```
108 pub fn set_default_cjk(is_cjk: bool) {
109 IS_CJK.store(is_cjk, Ordering::Relaxed);
110 }
111
112 #[deprecated(note = "Please use `char_opt` instead.")]
113 pub fn char(&self, ch: char) -> Option<usize> {
114 self.char_opt(ch)
115 }
116
117 /// Return the column width of a character.
118 ///
119 /// Return `None` for control characters
120 /// or other characters without a defined width.
121 ///
122 /// This is a wrapper of [`UnicodeWidthChar`].
123 /// It calls `width` or `width_cjk` depending on the configuration.
124 ///
125 /// [`UnicodeWidthChar`]: https://docs.rs/unicode-width/latest/unicode_width/trait.UnicodeWidthChar.html
126 ///
127 /// # Examples
128 /// ```
129 /// use unicode_width_utils::UnicodeWidth;
130 ///
131 /// let uw = UnicodeWidth::with_cjk(false);
132 /// assert_eq!(uw.char('A'), Some(1));
133 /// assert_eq!(uw.char('\n'), None);
134 /// ```
135 pub fn char_opt(&self, ch: char) -> Option<usize> {
136 match self.is_cjk {
137 false => UnicodeWidthChar::width(ch),
138 true => UnicodeWidthChar::width_cjk(ch),
139 }
140 }
141
142 /// Return the total column width of a string.
143 ///
144 /// This is a wrapper of [`UnicodeWidthStr`].
145 /// It calls `width` or `width_cjk` depending on the configuration.
146 ///
147 /// [`UnicodeWidthStr`]: https://docs.rs/unicode-width/latest/unicode_width/trait.UnicodeWidthStr.html
148 ///
149 /// # Examples
150 /// ```
151 /// use unicode_width_utils::UnicodeWidth;
152 ///
153 /// let uw = UnicodeWidth::with_cjk(false);
154 /// assert_eq!(uw.str("Hello"), 5);
155 /// ```
156 pub fn str(&self, str: &str) -> usize {
157 match self.is_cjk {
158 false => UnicodeWidthStr::width(str),
159 true => UnicodeWidthStr::width_cjk(str),
160 }
161 }
162
163 /// Set the tab size for [`truncate`].
164 /// Initially `0`.
165 ///
166 /// See also [`set_expand_tab`].
167 ///
168 /// [`set_expand_tab`]: UnicodeWidth::set_expand_tab
169 /// [`truncate`]: UnicodeWidth::truncate
170 ///
171 /// # Examples
172 /// ```
173 /// # use std::borrow::Cow;
174 /// use unicode_width_utils::UnicodeWidth;
175 ///
176 /// let mut uw = UnicodeWidth::new();
177 /// uw.set_tab_size(4);
178 /// assert_eq!(uw.truncate("A\tB", 3), Cow::Borrowed("A"));
179 /// assert_eq!(uw.truncate("A\tB", 4), Cow::Borrowed("A\t"));
180 /// assert_eq!(uw.truncate("A\tB", 5), Cow::Borrowed("A\tB"));
181 /// ```
182 pub fn set_tab_size(&mut self, tab_size: u8) {
183 self.tab_size = tab_size;
184 }
185
186 /// Set whether tabs should be expanded to spaces in [`truncate`].
187 /// Initially `false`.
188 ///
189 /// [`truncate`]: UnicodeWidth::truncate
190 ///
191 /// # Examples
192 /// ```
193 /// # use std::borrow::Cow;
194 /// use unicode_width_utils::UnicodeWidth;
195 ///
196 /// let mut uw = UnicodeWidth::new();
197 /// uw.set_tab_size(4);
198 /// uw.set_expand_tab(true);
199 /// assert_eq!(uw.truncate("A\tB", 3), Cow::Borrowed("A"));
200 /// assert_eq!(uw.truncate("A\tB", 4), Cow::Owned::<str>("A ".into()));
201 /// assert_eq!(uw.truncate("A\tB", 5), Cow::Owned::<str>("A B".into()));
202 /// ```
203 pub fn set_expand_tab(&mut self, should_expand_tab: bool) {
204 self.should_expand_tab = should_expand_tab;
205 }
206
207 /// Truncate a string slice to a maximum column width.
208 ///
209 /// The returned slice will be the longest prefix of `str`
210 /// whose total column width does not exceed `max_width`.
211 ///
212 /// See also [`set_tab_size`] and [`set_expand_tab`].
213 ///
214 /// [`set_tab_size`]: UnicodeWidth::set_tab_size
215 /// [`set_expand_tab`]: UnicodeWidth::set_expand_tab
216 ///
217 /// # Examples
218 /// ```
219 /// use unicode_width_utils::UnicodeWidth;
220 ///
221 /// let uw = UnicodeWidth::with_cjk(false);
222 /// assert_eq!(uw.truncate("hello", 3), "hel");
223 ///
224 /// // Truncating CJK text (each 'あ' is 2 columns wide).
225 /// let cjk = UnicodeWidth::with_cjk(true);
226 /// assert_eq!(cjk.truncate("あああ", 3), "あ");
227 /// assert_eq!(cjk.truncate("A█B", 2), "A");
228 /// ```
229 pub fn truncate<'a>(&self, input: &'a str, max_width: usize) -> Cow<'a, str> {
230 let mut width = 0;
231 let mut buffer: Option<String> = None;
232 let tab_size = self.tab_size as usize;
233 for (index, ch) in input.char_indices() {
234 let ch_width = if let Some(ch_width) = self.char_opt(ch) {
235 ch_width
236 } else if ch == '\t' && tab_size > 0 {
237 if buffer.is_none() && self.should_expand_tab {
238 let mut output = String::with_capacity(input.len() + tab_size * 4);
239 output.push_str(&input[..index]);
240 buffer = Some(output);
241 }
242 tab_size - (width % tab_size)
243 } else {
244 0
245 };
246 width += ch_width;
247 if width > max_width {
248 return match buffer {
249 None => Cow::Borrowed(&input[..index]),
250 Some(output) => Cow::Owned(output),
251 };
252 }
253
254 if let Some(ref mut output) = buffer {
255 if ch == '\t' {
256 for _ in 0..ch_width {
257 output.push(' ');
258 }
259 } else {
260 output.push(ch);
261 }
262 }
263 }
264 match buffer {
265 None => Cow::Borrowed(input),
266 Some(output) => Cow::Owned(output),
267 }
268 }
269}
270
271#[cfg(test)]
272mod tests {
273 use super::*;
274
275 #[test]
276 fn char() {
277 let uw = UnicodeWidth::with_cjk(false);
278 let cjk = UnicodeWidth::with_cjk(true);
279 assert_eq!(uw.char_opt('A'), Some(1));
280 assert_eq!(cjk.char_opt('A'), Some(1));
281 assert_eq!(uw.char_opt('\u{2588}'), Some(1));
282 assert_eq!(cjk.char_opt('\u{2588}'), Some(2));
283 assert_eq!(uw.char_opt('\u{3042}'), Some(2));
284 assert_eq!(cjk.char_opt('\u{3042}'), Some(2));
285 }
286
287 #[test]
288 fn str() {
289 let uw = UnicodeWidth::with_cjk(false);
290 let cjk = UnicodeWidth::with_cjk(true);
291 assert_eq!(uw.str("A"), 1);
292 assert_eq!(cjk.str("A"), 1);
293 assert_eq!(uw.str("\u{2588}"), 1);
294 assert_eq!(cjk.str("\u{2588}"), 2);
295 assert_eq!(uw.str("\u{3042}"), 2);
296 assert_eq!(cjk.str("\u{3042}"), 2);
297 }
298
299 #[test]
300 fn truncate() {
301 let uw = UnicodeWidth::with_cjk(false);
302 let cjk = UnicodeWidth::with_cjk(true);
303
304 assert_eq!(uw.truncate("hello", 0), "");
305 assert_eq!(uw.truncate("hello", 4), "hell");
306 assert_eq!(uw.truncate("hello", 5), "hello");
307 assert_eq!(uw.truncate("hello", 6), "hello");
308 assert_eq!(uw.truncate("hello", 10), "hello");
309
310 // \u{2588} is 1 column wide for `uw`, and 2 columns wide for `cjk`.
311 assert_eq!(uw.truncate("A\u{2588}B", 2), "A\u{2588}");
312 assert_eq!(cjk.truncate("A\u{2588}B", 2), "A");
313
314 // \u{3042} ('あ') is 2 columns wide.
315 assert_eq!(uw.truncate("\u{3042}", 1), "");
316 assert_eq!(uw.truncate("\u{3042}", 2), "\u{3042}");
317 assert_eq!(uw.truncate("\u{3042}", 3), "\u{3042}");
318
319 // Control characters with 0 width.
320 assert_eq!(uw.truncate("A\nB", 1), "A\n");
321 assert_eq!(uw.truncate("A\nB", 2), "A\nB");
322 assert_eq!(uw.truncate("\nA", 0), "\n");
323 }
324
325 #[test]
326 fn default_cjk() {
327 let original = IS_CJK.load(Ordering::Relaxed);
328
329 UnicodeWidth::set_default_cjk(false);
330 assert_eq!(UnicodeWidth::default().char_opt('\u{2588}'), Some(1));
331 assert_eq!(UnicodeWidth::new().char_opt('\u{2588}'), Some(1));
332
333 UnicodeWidth::set_default_cjk(true);
334 assert_eq!(UnicodeWidth::default().char_opt('\u{2588}'), Some(2));
335 assert_eq!(UnicodeWidth::new().char_opt('\u{2588}'), Some(2));
336
337 UnicodeWidth::set_default_cjk(original);
338 }
339
340 // If `tab_size` = 0, tab behaves as 0-width control character.
341 #[test]
342 fn truncate_tabs_zero() {
343 let mut uw = UnicodeWidth::new();
344 for _ in 0..2 {
345 assert_eq!(uw.truncate("A\tB", 1), Cow::Borrowed("A\t"));
346 assert_eq!(uw.truncate("A\tB", 2), Cow::Borrowed("A\tB"));
347 uw.set_expand_tab(true);
348 }
349 }
350
351 #[test]
352 fn truncate_tabs_expand_multi() {
353 let mut uw = UnicodeWidth::new();
354 uw.set_tab_size(4);
355 uw.set_expand_tab(true);
356 assert_eq!(uw.truncate("\t\t", 7), Cow::Owned::<str>(" ".into()));
357 assert_eq!(uw.truncate("\t\t", 8), Cow::Owned::<str>(" ".into()));
358 }
359}