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