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 /// # Examples
68 /// ```
69 /// use unicode_width_utils::UnicodeWidth;
70 ///
71 /// let non_cjk = UnicodeWidth::with_cjk(false);
72 /// assert_eq!(non_cjk.char('█'), 1);
73 ///
74 /// let cjk = UnicodeWidth::with_cjk(true);
75 /// assert_eq!(cjk.char('█'), 2);
76 /// ```
77 pub fn with_cjk(is_cjk: bool) -> Self {
78 Self {
79 is_cjk,
80 ..Default::default()
81 }
82 }
83
84 /// Set the default CJK configuration dynamically.
85 ///
86 /// Future instances
87 /// created using [`UnicodeWidth::new()`] or [`UnicodeWidth::default()`]
88 /// will inherit this default value
89 /// unless explicitly overridden with [`with_cjk()`].
90 ///
91 /// [`with_cjk()`]: `UnicodeWidth::with_cjk`
92 /// # Examples
93 /// ```
94 /// use unicode_width_utils::UnicodeWidth;
95 ///
96 /// // Set default CJK mode to true.
97 /// UnicodeWidth::set_default_cjk(true);
98 /// let uw = UnicodeWidth::new();
99 /// assert_eq!(uw.char('█'), 2);
100 ///
101 /// // Set default CJK mode to false.
102 /// UnicodeWidth::set_default_cjk(false);
103 /// let uw2 = UnicodeWidth::new();
104 /// assert_eq!(uw2.char('█'), 1);
105 /// ```
106 pub fn set_default_cjk(is_cjk: bool) {
107 IS_CJK.store(is_cjk, Ordering::Relaxed);
108 }
109
110 /// Set the size of control characters.
111 ///
112 /// # Examples
113 /// ```
114 /// use unicode_width_utils::UnicodeWidth;
115 ///
116 /// let mut uw = UnicodeWidth::new();
117 /// assert_eq!(uw.char('\t'), 1);
118 /// assert_eq!(uw.str("A\tB"), 3);
119 /// uw.set_control_size(0);
120 /// assert_eq!(uw.char('\t'), 0);
121 /// assert_eq!(uw.str("A\tB"), 2);
122 /// ```
123 pub fn set_control_size(&mut self, size: u8) {
124 self.control_size = size;
125 }
126
127 /// Return the column width of a character.
128 ///
129 /// This is a wrapper of [`UnicodeWidthChar`].
130 /// It calls `width` or `width_cjk` depending on the configuration.
131 ///
132 /// Control characters return `1` by default
133 /// to match [`UnicodeWidthStr`].
134 /// See also [`char_opt()`] and [`set_control_size()`] for control characters.
135 ///
136 /// [`char_opt()`]: UnicodeWidth::char_opt
137 /// [`set_control_size()`]: UnicodeWidth::set_control_size
138 ///
139 /// # Examples
140 /// ```
141 /// use unicode_width_utils::UnicodeWidth;
142 ///
143 /// let uw = UnicodeWidth::new();
144 /// assert_eq!(uw.char('A'), 1);
145 /// assert_eq!(uw.char('あ'), 2);
146 /// ```
147 pub fn char(&self, ch: char) -> usize {
148 self.char_opt(ch).unwrap_or(self.control_size as usize)
149 }
150
151 /// Return the column width of a character.
152 ///
153 /// Return `None` for control characters
154 /// or other characters without a defined width.
155 ///
156 /// This is a wrapper of [`UnicodeWidthChar`].
157 /// It calls `width` or `width_cjk` depending on the configuration.
158 ///
159 /// [`UnicodeWidthChar`]: https://docs.rs/unicode-width/latest/unicode_width/trait.UnicodeWidthChar.html
160 ///
161 /// # Examples
162 /// ```
163 /// use unicode_width_utils::UnicodeWidth;
164 ///
165 /// let uw = UnicodeWidth::new();
166 /// assert_eq!(uw.char_opt('A'), Some(1));
167 /// assert_eq!(uw.char_opt('\n'), None);
168 /// ```
169 pub fn char_opt(&self, ch: char) -> Option<usize> {
170 match self.is_cjk {
171 false => UnicodeWidthChar::width(ch),
172 true => UnicodeWidthChar::width_cjk(ch),
173 }
174 }
175
176 /// Return the total column width of a string.
177 ///
178 /// This is a wrapper of [`UnicodeWidthStr`].
179 /// It calls `width` or `width_cjk` depending on the configuration,
180 /// unless the [tab size is set][`set_tab_size()`]
181 /// or the [control character size is set][`set_control_size()`],
182 /// in which case the internal logic computes the width
183 /// by calling [`char_opt()`] or [`char()`] repeatedly.
184 ///
185 /// [`char()`]: UnicodeWidth::char
186 /// [`char_opt()`]: UnicodeWidth::char_opt
187 /// [`set_control_size()`]: UnicodeWidth::set_control_size
188 /// [`set_tab_size()`]: UnicodeWidth::set_tab_size
189 /// [`truncate()`]: UnicodeWidth::truncate
190 /// [`UnicodeWidthStr`]: https://docs.rs/unicode-width/latest/unicode_width/trait.UnicodeWidthStr.html
191 ///
192 /// # Examples
193 /// ```
194 /// use unicode_width_utils::UnicodeWidth;
195 ///
196 /// let mut uw = UnicodeWidth::new();
197 /// assert_eq!(uw.str("Hello"), 5);
198 /// assert_eq!(uw.str("Hello\t"), 6);
199 /// uw.set_tab_size(4);
200 /// assert_eq!(uw.str("Hello\t"), 8);
201 /// ```
202 pub fn str(&self, str: &str) -> usize {
203 if self.tab_size > 0 {
204 let mut iter = WidthIterator::new(self, str);
205 iter.consume_all();
206 return iter.width();
207 }
208 if self.control_size != 1 {
209 return str.chars().map(|ch| self.char(ch)).sum();
210 }
211 match self.is_cjk {
212 false => UnicodeWidthStr::width(str),
213 true => UnicodeWidthStr::width_cjk(str),
214 }
215 }
216
217 /// Set the tab size.
218 /// Initially `0`.
219 ///
220 /// This setting is used by [`truncate()`] and [`UnicodeWidth::str()`].
221 ///
222 /// See also [`set_expand_tab()`].
223 ///
224 /// [`set_expand_tab()`]: UnicodeWidth::set_expand_tab
225 /// [`truncate()`]: UnicodeWidth::truncate
226 ///
227 /// # Examples
228 /// ```
229 /// # use std::borrow::Cow;
230 /// use unicode_width_utils::UnicodeWidth;
231 ///
232 /// let mut uw = UnicodeWidth::new();
233 /// uw.set_tab_size(4);
234 /// assert_eq!(uw.truncate("A\tB", 3), Cow::Borrowed("A"));
235 /// assert_eq!(uw.truncate("A\tB", 4), Cow::Borrowed("A\t"));
236 /// assert_eq!(uw.truncate("A\tB", 5), Cow::Borrowed("A\tB"));
237 /// ```
238 pub fn set_tab_size(&mut self, tab_size: u8) {
239 self.tab_size = tab_size;
240 }
241
242 /// Set whether tabs should be expanded to spaces.
243 /// Initially `false`.
244 ///
245 /// This setting is used by [`truncate()`] and [`UnicodeWidth::str()`].
246 ///
247 /// See also [`set_tab_size()`].
248 ///
249 /// [`set_tab_size()`]: UnicodeWidth::set_tab_size
250 /// [`truncate()`]: UnicodeWidth::truncate
251 ///
252 /// # Examples
253 /// ```
254 /// # use std::borrow::Cow;
255 /// use unicode_width_utils::UnicodeWidth;
256 ///
257 /// let mut uw = UnicodeWidth::new();
258 /// uw.set_tab_size(4);
259 /// uw.set_expand_tab(true);
260 /// assert_eq!(uw.truncate("A\tB", 3), Cow::Borrowed("A"));
261 /// assert_eq!(uw.truncate("A\tB", 4), Cow::Owned::<str>("A ".into()));
262 /// assert_eq!(uw.truncate("A\tB", 5), Cow::Owned::<str>("A B".into()));
263 /// ```
264 pub fn set_expand_tab(&mut self, should_expand_tab: bool) {
265 self.should_expand_tab = should_expand_tab;
266 }
267
268 /// Truncate a string slice to a maximum column width.
269 ///
270 /// The returned slice will be the longest prefix of `input`
271 /// whose total column width does not exceed `max_width`.
272 ///
273 /// See also [`set_tab_size()`] and [`set_expand_tab()`].
274 ///
275 /// [`set_tab_size()`]: UnicodeWidth::set_tab_size
276 /// [`set_expand_tab()`]: UnicodeWidth::set_expand_tab
277 ///
278 /// # Examples
279 /// ```
280 /// use unicode_width_utils::UnicodeWidth;
281 ///
282 /// let uw = UnicodeWidth::new();
283 /// assert_eq!(uw.truncate("hello", 3), "hel");
284 ///
285 /// // Truncating CJK text (each 'あ' is 2 columns wide).
286 /// assert_eq!(uw.truncate("あああ", 3), "あ");
287 /// assert_eq!(uw.truncate("A█B", 2), "A█");
288 ///
289 /// let cjk = UnicodeWidth::with_cjk(true);
290 /// assert_eq!(cjk.truncate("A█B", 2), "A");
291 /// ```
292 pub fn truncate<'a>(&self, input: &'a str, max_width: usize) -> Cow<'a, str> {
293 let mut iter = WidthIterator::new(self, input);
294 iter.set_max_width(max_width);
295 iter.consume_all();
296 iter.into()
297 }
298}
299
300#[cfg(test)]
301mod tests {
302 use super::*;
303
304 #[test]
305 fn char() {
306 let uw = UnicodeWidth::with_cjk(false);
307 let cjk = UnicodeWidth::with_cjk(true);
308 assert_eq!(uw.char('A'), 1);
309 assert_eq!(cjk.char('A'), 1);
310 assert_eq!(uw.char('\u{2588}'), 1);
311 assert_eq!(cjk.char('\u{2588}'), 2);
312 assert_eq!(uw.char('\u{3042}'), 2);
313 assert_eq!(cjk.char('\u{3042}'), 2);
314 }
315
316 #[test]
317 fn str() {
318 let uw = UnicodeWidth::with_cjk(false);
319 let cjk = UnicodeWidth::with_cjk(true);
320 assert_eq!(uw.str("A"), 1);
321 assert_eq!(cjk.str("A"), 1);
322 assert_eq!(uw.str("\u{2588}"), 1);
323 assert_eq!(cjk.str("\u{2588}"), 2);
324 assert_eq!(uw.str("\u{3042}"), 2);
325 assert_eq!(cjk.str("\u{3042}"), 2);
326 }
327
328 #[test]
329 fn str_tab() {
330 let mut uw = UnicodeWidth::with_cjk(false);
331 uw.set_tab_size(4);
332 assert_eq!(uw.str("A"), 1);
333 assert_eq!(uw.str("A\t"), 4);
334 assert_eq!(uw.str("A\tB"), 5);
335 assert_eq!(uw.str("A\t\tB"), 9);
336 }
337
338 #[test]
339 fn truncate() {
340 let uw = UnicodeWidth::with_cjk(false);
341 let cjk = UnicodeWidth::with_cjk(true);
342
343 assert_eq!(uw.truncate("hello", 0), "");
344 assert_eq!(uw.truncate("hello", 4), "hell");
345 assert_eq!(uw.truncate("hello", 5), "hello");
346 assert_eq!(uw.truncate("hello", 6), "hello");
347 assert_eq!(uw.truncate("hello", 10), "hello");
348
349 // \u{2588} is 1 column wide for `uw`, and 2 columns wide for `cjk`.
350 assert_eq!(uw.truncate("A\u{2588}B", 2), "A\u{2588}");
351 assert_eq!(cjk.truncate("A\u{2588}B", 2), "A");
352
353 // \u{3042} ('あ') is 2 columns wide.
354 assert_eq!(uw.truncate("\u{3042}", 1), "");
355 assert_eq!(uw.truncate("\u{3042}", 2), "\u{3042}");
356 assert_eq!(uw.truncate("\u{3042}", 3), "\u{3042}");
357
358 // Control characters with 1 column wide.
359 assert_eq!(uw.truncate("A\nB", 1), "A");
360 assert_eq!(uw.truncate("A\nB", 2), "A\n");
361 assert_eq!(uw.truncate("\nA", 0), "");
362 assert_eq!(uw.truncate("\nA", 1), "\n");
363 }
364
365 #[test]
366 fn default_cjk() {
367 let original = IS_CJK.load(Ordering::Relaxed);
368
369 UnicodeWidth::set_default_cjk(false);
370 assert_eq!(UnicodeWidth::default().char('\u{2588}'), 1);
371 assert_eq!(UnicodeWidth::new().char('\u{2588}'), 1);
372
373 UnicodeWidth::set_default_cjk(true);
374 assert_eq!(UnicodeWidth::default().char('\u{2588}'), 2);
375 assert_eq!(UnicodeWidth::new().char('\u{2588}'), 2);
376
377 UnicodeWidth::set_default_cjk(original);
378 }
379
380 // If `tab_size` = 0, tab behaves as 1 column wide control character.
381 #[test]
382 fn truncate_tabs_no_tab_size() {
383 let mut uw = UnicodeWidth::new();
384 for _ in 0..2 {
385 assert_eq!(uw.truncate("A\tB", 1), Cow::Borrowed("A"));
386 assert_eq!(uw.truncate("A\tB", 2), Cow::Borrowed("A\t"));
387 assert_eq!(uw.truncate("A\tB", 3), Cow::Borrowed("A\tB"));
388 uw.set_expand_tab(true);
389 }
390 }
391
392 #[test]
393 fn truncate_tabs_expand_multi() {
394 let mut uw = UnicodeWidth::new();
395 uw.set_tab_size(4);
396 uw.set_expand_tab(true);
397 assert_eq!(uw.truncate("\t\t", 7), Cow::Owned::<str>(" ".into()));
398 assert_eq!(uw.truncate("\t\t", 8), Cow::Owned::<str>(" ".into()));
399 }
400}