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