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
6use crate::LineIterator;
7use crate::WidthIterator;
8
9static IS_CJK: LazyLock<AtomicBool> = LazyLock::new(|| {
10 let is_cjk = match std::env::var("UNICODE_WIDTH") {
11 Ok(value) => value.eq_ignore_ascii_case("cjk"),
12 _ => false,
13 };
14 AtomicBool::new(is_cjk)
15});
16
17/// A configuration helper to measure character and string widths.
18///
19/// It determines the width of Unicode characters and strings,
20/// optionally treating East Asian Ambiguous characters
21/// (such as certain Greek, Cyrillic, and CJK characters)
22/// as having a width of 2 (CJK mode).
23///
24/// The default CJK mode is initialized at startup
25/// based on the `UNICODE_WIDTH` environment variable,
26/// but can also be dynamically modified using [`set_default_cjk()`].
27///
28/// [`set_default_cjk()`]: UnicodeWidth::set_default_cjk
29#[derive(Clone, Debug)]
30pub struct UnicodeWidth {
31 is_cjk: bool,
32 pub(crate) should_expand_tab: bool,
33 pub(crate) tab_size: u8,
34 pub(crate) control_size: u8,
35}
36
37impl Default for UnicodeWidth {
38 fn default() -> Self {
39 Self {
40 is_cjk: IS_CJK.load(Ordering::Relaxed),
41 should_expand_tab: false,
42 tab_size: 0,
43 control_size: 1,
44 }
45 }
46}
47
48impl UnicodeWidth {
49 /// Create a `UnicodeWidth` instance using the default CJK mode.
50 ///
51 /// The default CJK mode is determined by the global setting,
52 /// which defaults to the value of the `UNICODE_WIDTH` environment variable
53 /// (value `"cjk"` enabling CJK mode).
54 ///
55 /// See also [`set_default_cjk()`].
56 ///
57 /// [`set_default_cjk()`]: UnicodeWidth::set_default_cjk
58 pub fn new() -> Self {
59 Self::default()
60 }
61
62 /// Create a new `UnicodeWidth` instance with a specific CJK flag.
63 ///
64 /// If `is_cjk` is true,
65 /// East Asian Ambiguous characters will be treated as 2 columns wide.
66 /// If false, they will be treated as 1 column wide.
67 ///
68 /// See also [`set_cjk()`].
69 ///
70 /// [`set_cjk()`]: UnicodeWidth::set_cjk
71 ///
72 /// # Examples
73 /// ```
74 /// use unicode_width_utils::UnicodeWidth;
75 ///
76 /// let non_cjk = UnicodeWidth::with_cjk(false);
77 /// assert_eq!(non_cjk.char('█'), 1);
78 /// let cjk = UnicodeWidth::with_cjk(true);
79 /// assert_eq!(cjk.char('█'), 2);
80 /// ```
81 pub fn with_cjk(is_cjk: bool) -> Self {
82 Self {
83 is_cjk,
84 ..Default::default()
85 }
86 }
87
88 /// Set to whether to perform an alternate width calculation
89 /// more suited to CJK contexts or not.
90 ///
91 /// When set to `true`,
92 /// characters in the Ambiguous category according to
93 /// [Unicode Standard Annex #11] as 2 columns wide.
94 ///
95 /// See also the ["cjk" feature flag] and
96 /// [`UnicodeWidthChar::width_cjk()`].
97 ///
98 /// ["cjk" feature flag]: https://docs.rs/unicode-width/latest/unicode_width/#cjk-feature-flag
99 /// [Unicode Standard Annex #11]: https://www.unicode.org/reports/tr11/
100 ///
101 /// # Examples
102 /// ```
103 /// use unicode_width_utils::UnicodeWidth;
104 ///
105 /// let mut uw = UnicodeWidth::with_cjk(false);
106 /// assert_eq!(uw.char('█'), 1);
107 /// uw.set_cjk(true);
108 /// assert_eq!(uw.char('█'), 2);
109 /// ```
110 #[inline]
111 pub fn set_cjk(&mut self, is_cjk: bool) {
112 self.is_cjk = is_cjk;
113 }
114
115 /// Set the default CJK configuration dynamically.
116 ///
117 /// Future instances
118 /// created using [`UnicodeWidth::new()`] or [`UnicodeWidth::default()`]
119 /// will inherit this default value
120 /// unless explicitly overridden with [`with_cjk()`] or [`set_cjk()`].
121 ///
122 /// [`set_cjk()`]: UnicodeWidth::set_cjk
123 /// [`with_cjk()`]: `UnicodeWidth::with_cjk`
124 /// # Examples
125 /// ```
126 /// use unicode_width_utils::UnicodeWidth;
127 ///
128 /// // Set default CJK mode to true.
129 /// UnicodeWidth::set_default_cjk(true);
130 /// let cjk = UnicodeWidth::new();
131 /// assert_eq!(cjk.char('█'), 2);
132 ///
133 /// // Set default CJK mode to false.
134 /// UnicodeWidth::set_default_cjk(false);
135 /// let non_cjk = UnicodeWidth::new();
136 /// assert_eq!(non_cjk.char('█'), 1);
137 /// ```
138 pub fn set_default_cjk(is_cjk: bool) {
139 IS_CJK.store(is_cjk, Ordering::Relaxed);
140 }
141
142 /// Set the size of control characters.
143 ///
144 /// # Examples
145 /// ```
146 /// use unicode_width_utils::UnicodeWidth;
147 ///
148 /// let mut uw = UnicodeWidth::new();
149 /// assert_eq!(uw.char('\t'), 1);
150 /// assert_eq!(uw.str("A\tB"), 3);
151 /// uw.set_control_size(0);
152 /// assert_eq!(uw.char('\t'), 0);
153 /// assert_eq!(uw.str("A\tB"), 2);
154 /// ```
155 #[inline]
156 pub fn set_control_size(&mut self, size: u8) {
157 self.control_size = size;
158 }
159
160 /// Set the tab size.
161 /// Initially `0`.
162 ///
163 /// This setting is used by [`truncate()`] and [`UnicodeWidth::str()`].
164 ///
165 /// See also [`set_expand_tab()`].
166 ///
167 /// [`set_expand_tab()`]: UnicodeWidth::set_expand_tab
168 /// [`truncate()`]: UnicodeWidth::truncate
169 ///
170 /// # Examples
171 /// ```
172 /// # use std::borrow::Cow;
173 /// use unicode_width_utils::UnicodeWidth;
174 ///
175 /// let mut uw = UnicodeWidth::new();
176 /// uw.set_tab_size(4);
177 /// assert_eq!(uw.truncate("A\tB", 3), Cow::Borrowed("A"));
178 /// assert_eq!(uw.truncate("A\tB", 4), Cow::Borrowed("A\t"));
179 /// assert_eq!(uw.truncate("A\tB", 5), Cow::Borrowed("A\tB"));
180 /// ```
181 #[inline]
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.
187 /// Initially `false`.
188 ///
189 /// This setting is used by [`truncate()`] and [`UnicodeWidth::str()`].
190 ///
191 /// See also [`set_tab_size()`].
192 ///
193 /// [`set_tab_size()`]: UnicodeWidth::set_tab_size
194 /// [`truncate()`]: UnicodeWidth::truncate
195 ///
196 /// # Examples
197 /// ```
198 /// # use std::borrow::Cow;
199 /// use unicode_width_utils::UnicodeWidth;
200 ///
201 /// let mut uw = UnicodeWidth::new();
202 /// uw.set_tab_size(4);
203 /// uw.set_expand_tab(true);
204 /// assert_eq!(uw.truncate("A\tB", 3), Cow::Borrowed("A"));
205 /// assert_eq!(uw.truncate("A\tB", 4), Cow::Owned::<str>("A ".into()));
206 /// assert_eq!(uw.truncate("A\tB", 5), Cow::Owned::<str>("A B".into()));
207 /// ```
208 #[inline]
209 pub fn set_expand_tab(&mut self, should_expand_tab: bool) {
210 self.should_expand_tab = should_expand_tab;
211 }
212
213 /// Return the column width of a character.
214 ///
215 /// This is a wrapper of [`UnicodeWidthChar`].
216 /// It calls `width` or `width_cjk` depending on the configuration.
217 ///
218 /// Control characters return `1` by default
219 /// to match [`UnicodeWidthStr`].
220 /// See also [`char_opt()`] and [`set_control_size()`] for control characters.
221 ///
222 /// [`char_opt()`]: UnicodeWidth::char_opt
223 /// [`set_control_size()`]: UnicodeWidth::set_control_size
224 ///
225 /// # Examples
226 /// ```
227 /// use unicode_width_utils::UnicodeWidth;
228 ///
229 /// let uw = UnicodeWidth::new();
230 /// assert_eq!(uw.char('A'), 1);
231 /// assert_eq!(uw.char('あ'), 2);
232 /// ```
233 pub fn char(&self, ch: char) -> usize {
234 self.char_opt(ch).unwrap_or(self.control_size as usize)
235 }
236
237 /// Return the column width of a character.
238 ///
239 /// Return `None` for control characters
240 /// or other characters without a defined width.
241 ///
242 /// This is a wrapper of [`UnicodeWidthChar`].
243 /// It calls `width` or `width_cjk` depending on the configuration.
244 ///
245 /// [`UnicodeWidthChar`]: https://docs.rs/unicode-width/latest/unicode_width/trait.UnicodeWidthChar.html
246 ///
247 /// # Examples
248 /// ```
249 /// use unicode_width_utils::UnicodeWidth;
250 ///
251 /// let uw = UnicodeWidth::new();
252 /// assert_eq!(uw.char_opt('A'), Some(1));
253 /// assert_eq!(uw.char_opt('\n'), None);
254 /// ```
255 pub fn char_opt(&self, ch: char) -> Option<usize> {
256 match self.is_cjk {
257 false => UnicodeWidthChar::width(ch),
258 true => UnicodeWidthChar::width_cjk(ch),
259 }
260 }
261
262 /// Return the total column width of a string.
263 ///
264 /// This is a wrapper of [`UnicodeWidthStr`].
265 /// It calls `width` or `width_cjk` depending on the configuration,
266 /// unless the [tab size is set][`set_tab_size()`]
267 /// or the [control character size is set][`set_control_size()`],
268 /// in which case the internal logic computes the width
269 /// by calling [`char_opt()`] or [`char()`] repeatedly.
270 ///
271 /// [`char()`]: UnicodeWidth::char
272 /// [`char_opt()`]: UnicodeWidth::char_opt
273 /// [`set_control_size()`]: UnicodeWidth::set_control_size
274 /// [`set_tab_size()`]: UnicodeWidth::set_tab_size
275 /// [`truncate()`]: UnicodeWidth::truncate
276 /// [`UnicodeWidthStr`]: https://docs.rs/unicode-width/latest/unicode_width/trait.UnicodeWidthStr.html
277 ///
278 /// # Examples
279 /// ```
280 /// use unicode_width_utils::UnicodeWidth;
281 ///
282 /// let mut uw = UnicodeWidth::new();
283 /// assert_eq!(uw.str("Hello"), 5);
284 /// assert_eq!(uw.str("Hello\t"), 6);
285 /// uw.set_tab_size(4);
286 /// assert_eq!(uw.str("Hello\t"), 8);
287 /// ```
288 pub fn str(&self, str: &str) -> usize {
289 if self.tab_size > 0 {
290 let mut iter = WidthIterator::new(self, str);
291 iter.consume_all();
292 return iter.width();
293 }
294 if self.control_size != 1 {
295 return str.chars().map(|ch| self.char(ch)).sum();
296 }
297 match self.is_cjk {
298 false => UnicodeWidthStr::width(str),
299 true => UnicodeWidthStr::width_cjk(str),
300 }
301 }
302
303 /// Truncate a string slice to a maximum column width.
304 ///
305 /// The returned slice will be the longest prefix of `input`
306 /// whose total column width does not exceed `max_width`.
307 ///
308 /// See also [`set_tab_size()`] and [`set_expand_tab()`].
309 ///
310 /// [`set_tab_size()`]: UnicodeWidth::set_tab_size
311 /// [`set_expand_tab()`]: UnicodeWidth::set_expand_tab
312 ///
313 /// # Examples
314 /// ```
315 /// use unicode_width_utils::UnicodeWidth;
316 ///
317 /// let uw = UnicodeWidth::new();
318 /// assert_eq!(uw.truncate("hello", 3), "hel");
319 ///
320 /// // Truncating CJK text (each 'あ' is 2 columns wide).
321 /// assert_eq!(uw.truncate("あああ", 3), "あ");
322 /// assert_eq!(uw.truncate("A█B", 2), "A█");
323 ///
324 /// let cjk = UnicodeWidth::with_cjk(true);
325 /// assert_eq!(cjk.truncate("A█B", 2), "A");
326 /// ```
327 pub fn truncate<'a>(&self, input: &'a str, max_width: usize) -> Cow<'a, str> {
328 let mut iter = WidthIterator::new(self, input);
329 iter.set_max_width(max_width);
330 iter.consume_all();
331 iter.into()
332 }
333
334 /// Return a [`LineIterator`]
335 /// to iterate over chunks of a string based on display width.
336 ///
337 /// Unlike [`truncate()`],
338 /// each line is guaranteed to have at least one character,
339 /// even if it is wider than `max_width`.
340 ///
341 /// [`truncate()`]: UnicodeWidth::truncate
342 ///
343 /// # Examples
344 /// ```
345 /// use unicode_width_utils::UnicodeWidth;
346 ///
347 /// let uw = UnicodeWidth::new();
348 /// assert_eq!(
349 /// uw.lines("12345678", 3).collect::<Vec<_>>(),
350 /// vec!["123", "456", "78"]
351 /// );
352 /// ```
353 pub fn lines<'a>(&self, input: &'a str, max_width: usize) -> LineIterator<'_, 'a> {
354 LineIterator::new(self, input, max_width)
355 }
356}
357
358#[cfg(test)]
359mod tests {
360 use super::*;
361
362 #[test]
363 fn char() {
364 let uw = UnicodeWidth::with_cjk(false);
365 let cjk = UnicodeWidth::with_cjk(true);
366 assert_eq!(uw.char('A'), 1);
367 assert_eq!(cjk.char('A'), 1);
368 assert_eq!(uw.char('\u{2588}'), 1);
369 assert_eq!(cjk.char('\u{2588}'), 2);
370 assert_eq!(uw.char('\u{3042}'), 2);
371 assert_eq!(cjk.char('\u{3042}'), 2);
372 }
373
374 #[test]
375 fn str() {
376 let uw = UnicodeWidth::with_cjk(false);
377 let cjk = UnicodeWidth::with_cjk(true);
378 assert_eq!(uw.str("A"), 1);
379 assert_eq!(cjk.str("A"), 1);
380 assert_eq!(uw.str("\u{2588}"), 1);
381 assert_eq!(cjk.str("\u{2588}"), 2);
382 assert_eq!(uw.str("\u{3042}"), 2);
383 assert_eq!(cjk.str("\u{3042}"), 2);
384 }
385
386 #[test]
387 fn str_tab() {
388 let mut uw = UnicodeWidth::with_cjk(false);
389 uw.set_tab_size(4);
390 assert_eq!(uw.str("A"), 1);
391 assert_eq!(uw.str("A\t"), 4);
392 assert_eq!(uw.str("A\tB"), 5);
393 assert_eq!(uw.str("A\t\tB"), 9);
394 }
395
396 #[test]
397 fn truncate() {
398 let uw = UnicodeWidth::with_cjk(false);
399 let cjk = UnicodeWidth::with_cjk(true);
400
401 assert_eq!(uw.truncate("hello", 0), "");
402 assert_eq!(uw.truncate("hello", 4), "hell");
403 assert_eq!(uw.truncate("hello", 5), "hello");
404 assert_eq!(uw.truncate("hello", 6), "hello");
405 assert_eq!(uw.truncate("hello", 10), "hello");
406
407 // \u{2588} is 1 column wide for `uw`, and 2 columns wide for `cjk`.
408 assert_eq!(uw.truncate("A\u{2588}B", 2), "A\u{2588}");
409 assert_eq!(cjk.truncate("A\u{2588}B", 2), "A");
410
411 // \u{3042} ('あ') is 2 columns wide.
412 assert_eq!(uw.truncate("\u{3042}", 1), "");
413 assert_eq!(uw.truncate("\u{3042}", 2), "\u{3042}");
414 assert_eq!(uw.truncate("\u{3042}", 3), "\u{3042}");
415
416 // Control characters with 1 column wide.
417 assert_eq!(uw.truncate("A\nB", 1), "A");
418 assert_eq!(uw.truncate("A\nB", 2), "A\n");
419 assert_eq!(uw.truncate("\nA", 0), "");
420 assert_eq!(uw.truncate("\nA", 1), "\n");
421 }
422
423 #[test]
424 fn default_cjk() {
425 let original = IS_CJK.load(Ordering::Relaxed);
426
427 UnicodeWidth::set_default_cjk(false);
428 assert_eq!(UnicodeWidth::default().char('\u{2588}'), 1);
429 assert_eq!(UnicodeWidth::new().char('\u{2588}'), 1);
430
431 UnicodeWidth::set_default_cjk(true);
432 assert_eq!(UnicodeWidth::default().char('\u{2588}'), 2);
433 assert_eq!(UnicodeWidth::new().char('\u{2588}'), 2);
434
435 UnicodeWidth::set_default_cjk(original);
436 }
437
438 // If `tab_size` = 0, tab behaves as 1 column wide control character.
439 #[test]
440 fn truncate_tabs_no_tab_size() {
441 let mut uw = UnicodeWidth::new();
442 for _ in 0..2 {
443 assert_eq!(uw.truncate("A\tB", 1), Cow::Borrowed("A"));
444 assert_eq!(uw.truncate("A\tB", 2), Cow::Borrowed("A\t"));
445 assert_eq!(uw.truncate("A\tB", 3), Cow::Borrowed("A\tB"));
446 uw.set_expand_tab(true);
447 }
448 }
449
450 #[test]
451 fn truncate_tabs_expand_multi() {
452 let mut uw = UnicodeWidth::new();
453 uw.set_tab_size(4);
454 uw.set_expand_tab(true);
455 assert_eq!(uw.truncate("\t\t", 7), Cow::Owned::<str>(" ".into()));
456 assert_eq!(uw.truncate("\t\t", 8), Cow::Owned::<str>(" ".into()));
457 }
458
459 #[test]
460 fn lines_tabs_expand() {
461 let mut uw = UnicodeWidth::new();
462 uw.set_tab_size(4);
463 uw.set_expand_tab(true);
464 let lines: Vec<_> = uw.lines("hi\tworld rust", 8).collect();
465 assert_eq!(lines, vec!["hi worl", "d rust"]);
466 }
467}