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