Fixed cell width issues when using ANSI color codes.

This commit adds a `utils::display_width` function, which is just a wrapper
around `UnicodeWidthStr::width` which also takes ANSI color codes into account.

This is required when creating cells from strings which are already colored
using ANSI color codes (instead of coloring the cells using styles). Since
color codes are of the form \u{1b}[ ... m, but UnicodeWidthStr::width only
takes the first \u{1b} into account, this would create cell width issues.

2 files changed, 33 insertions, 3 deletions

diff --git a/src/cell.rs b/src/cell.rs
index 255c3fd..cca22d4 100644
--- a/src/cell.rs
+++ b/src/cell.rs
@@ -2,9 +2,9 @@
 
 use std::io::{Write, Error};
 use std::string::ToString;
-use unicode_width::UnicodeWidthStr;
 use super::{Attr, Terminal, color};
 use super::format::Alignment;
+use super::utils::display_width;
 use super::utils::print_align;
 
 /// Represent a table cell containing a string.
@@ -26,7 +26,7 @@ impl Cell {
         let content: Vec<String> = string.lines().map(|x| x.to_string()).collect();
         let mut width = 0;
         for cont in &content {
-            let l = UnicodeWidthStr::width(&cont[..]);
+            let l = display_width(&cont[..]);
             if l > width {
                 width = l;
             }
diff --git a/src/utils.rs b/src/utils.rs
index b96828e..79e3be8 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -57,7 +57,7 @@ pub fn print_align<T: Write + ?Sized>(out: &mut T,
                                       size: usize,
                                       skip_right_fill: bool)
                                       -> Result<(), Error> {
-    let text_len = UnicodeWidthStr::width(text);
+    let text_len = display_width(text);
     let mut nfill = if text_len < size { size - text_len } else { 0 };
     let n = match align {
         Alignment::LEFT => 0,
@@ -75,6 +75,36 @@ pub fn print_align<T: Write + ?Sized>(out: &mut T,
     Ok(())
 }
 
+/// Return the display width of a unicode string.
+/// This functions takes ANSI-escaped color codes into account.
+pub fn display_width(text: &str) -> usize {
+    let width = UnicodeWidthStr::width(text);
+    let mut state = 0;
+    let mut hidden = 0;
+
+    for c in text.chars() {
+        state = match (state, c) {
+            (0, '\u{1b}') => 1,
+            (1, '[') => 2,
+            (1, _) => 0,
+            (2, 'm') => 3,
+            _ => state,
+        };
+
+        // We don't count escape characters as hidden as
+        // UnicodeWidthStr::width already considers them.
+        if state > 1 {
+            hidden += 1;
+        }
+
+        if state == 3 {
+            state = 0;
+        }
+    }
+
+    width - hidden
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;