handle ansi style in cell content

Cargo.toml

@@ -42,6 +42,7 @@ crossterm = { version = "0.25", optional = true }
 strum = "0.24"
 strum_macros = "0.24"
 unicode-width = "0.1"
+console = "0.15.1"
 
 [dev-dependencies]
 pretty_assertions = "1"

examples/inner_style.rs

@@ -0,0 +1,33 @@
+use comfy_table::{Cell, ContentArrangement, Row, Table};
+
+fn main() {
+    let mut table = Table::new();
+    //table.load_preset(comfy_table::presets::NOTHING);
+    table.set_content_arrangement(ContentArrangement::Dynamic);
+    table.set_width(85);
+
+    let mut row = Row::new();
+    row.add_cell(Cell::new(format!(
+        "hello{}cell1",
+        console::style("123\n456").dim().blue()
+    )));
+    row.add_cell(Cell::new("cell2"));
+
+    table.add_row(row);
+
+    let mut row = Row::new();
+    row.add_cell(Cell::new(
+        format!(r"cell sys-devices-pci00:00-0000:000:07:00.1-usb2-2\x2d1-2\x2d1.3-2\x2d1.3:1.0-host2-target2:0:0-2:0:0:1-block-sdb{}", console::style(".device").bold().red())
+    ));
+    row.add_cell(Cell::new(
+        "cell4 asdfasfsad asdfasdf sad fas df asdf as df asdf    asdfasdfasdfasdfasdfasdfa dsfa sdf asdf asd f asdf as df sadf asd fas df "
+    ));
+    table.add_row(row);
+
+    let mut row = Row::new();
+    row.add_cell(Cell::new("cell5"));
+    row.add_cell(Cell::new("cell6"));
+    table.add_row(row);
+
+    println!("{}", table);
+}

src/cell.rs

@@ -1,7 +1,7 @@
 #[cfg(feature = "tty")]
 use crossterm::style::{Attribute, Color};
 
-use crate::style::CellAlignment;
+use crate::{style::CellAlignment, utils::fix_style_in_split_str};
 
 /// A stylable table cell with content.
 #[derive(Clone, Debug)]
@@ -26,12 +26,14 @@ impl Cell {
     /// Create a new Cell
     #[allow(clippy::needless_pass_by_value)]
     pub fn new<T: ToString>(content: T) -> Self {
+        let content = content.to_string();
+        let mut split_content: Vec<String> = content.split('\n').map(ToString::to_string).collect();
+
+        // corrects ansi codes so style is terminated and resumed around the split
+        fix_style_in_split_str(&mut split_content);
+
         Self {
-            content: content
-                .to_string()
-                .split('\n')
-                .map(ToString::to_string)
-                .collect(),
+            content: split_content,
             delimiter: None,
             alignment: None,
             #[cfg(feature = "tty")]

src/row.rs

@@ -1,7 +1,5 @@
 use std::slice::Iter;
 
-use unicode_width::UnicodeWidthStr;
-
 use crate::cell::{Cell, Cells};
 
 /// Each row contains [Cells](crate::Cell) and can be added to a [Table](crate::Table).
@@ -63,7 +61,11 @@ impl Row {
                 // Each entry represents the longest string width for a cell.
                 cell.content
                     .iter()
-                    .map(|string| string.width())
+                    .map(|string| {
+                        //let width = console::measure_text_width(string);
+                        //println!("{} {} {}", width, &string, string.escape_debug());
+                        console::measure_text_width(string)
+                    })
                     .max()
                     .unwrap_or(0)
             })

src/utils/formatting/content_format.rs

@@ -1,3 +1,4 @@
+use console::strip_ansi_codes;
 #[cfg(feature = "tty")]
 use crossterm::style::{style, Stylize};
 use unicode_width::UnicodeWidthStr;
@@ -87,7 +88,7 @@ pub fn format_row(
         // Iterate over each line and split it into multiple lines, if necessary.
         // Newlines added by the user will be preserved.
         for line in cell.content.iter() {
-            if line.width() > info.content_width.into() {
+            if console::measure_text_width(line) > info.content_width.into() {
                 let mut splitted = split_line(line, info, delimiter);
                 cell_lines.append(&mut splitted);
             } else {
@@ -106,6 +107,11 @@ pub fn format_row(
                     .get_mut(lines - 1)
                     .expect("We know it's this long.");
 
+                // we are truncate the line, so we might cuttoff a ansi code
+                let stripped = strip_ansi_codes(last_line).to_string();
+                last_line.clear();
+                last_line.push_str(&stripped);
+
                 // Don't do anything if the collumn is smaller then 6 characters
                 let width: usize = info.content_width.into();
                 if width >= 6 {
@@ -177,7 +183,8 @@ pub fn format_row(
 #[allow(unused_variables)]
 fn align_line(table: &Table, info: &ColumnDisplayInfo, cell: &Cell, mut line: String) -> String {
     let content_width = info.content_width;
-    let remaining: usize = usize::from(content_width).saturating_sub(line.width());
+    let remaining: usize =
+        usize::from(content_width).saturating_sub(console::measure_text_width(&line));
 
     // Apply the styling before aligning the line, if the user requests it.
     // That way non-delimiter whitespaces won't have stuff like underlines.

src/utils/formatting/content_split.rs

@@ -28,8 +28,8 @@ pub fn split_line(line: &str, info: &ColumnDisplayInfo, delimiter: char) -> Vec<
 
     let mut current_line = String::new();
     while let Some(next) = elements.pop() {
-        let current_length = current_line.width();
-        let next_length = next.width();
+        let current_length = console::measure_text_width(&current_line);
+        let next_length = console::measure_text_width(&next);
 
         // Some helper variables
         // The length of the current line when combining it with the next element
@@ -39,7 +39,7 @@ pub fn split_line(line: &str, info: &ColumnDisplayInfo, delimiter: char) -> Vec<
             added_length += 1;
         }
         // The remaining width for this column. If we are on a non-empty line, subtract 1 for the delimiter.
-        let mut remaining_width = content_width - current_line.width();
+        let mut remaining_width = content_width - current_length;
         if !current_line.is_empty() {
             remaining_width = remaining_width.saturating_sub(1);
         }
@@ -148,45 +148,112 @@ const MIN_FREE_CHARS: usize = 2;
 /// Otherwise, we simply return the current line and basically don't do anything.
 fn check_if_full(lines: &mut Vec<String>, content_width: usize, current_line: String) -> String {
     // Already complete the current line, if there isn't space for more than two chars
-    if current_line.width() > content_width.saturating_sub(MIN_FREE_CHARS) {
+    if console::measure_text_width(&current_line) > content_width.saturating_sub(MIN_FREE_CHARS) {
         lines.push(current_line);
         return String::new();
     }
 
     current_line
 }
 
-/// Splits a long word at a given character width.
+const ANSI_RESET : &str = "\u{1b}[0m";
+
+/// Splits a long word at a given character width. Inserting the needed ansi codes to preserve style.
 /// This needs some special logic, as we have to take multi-character UTF-8 symbols into account.
 /// When simply splitting at a certain char position, we might end up with a string that's has a
 /// wider display width than allowed.
 fn split_long_word(allowed_width: usize, word: &str) -> (String, String) {
-    let mut current_width = 0;
-    let mut splitted = String::new();
-
-    let mut char_iter = word.chars().peekable();
-    // Check if the string might be too long, one character at a time.
-    // Peek into the next char and check the exit condition.
-    // That is, pushing the next character would result in the string being too long.
-    while let Some(c) = char_iter.peek() {
-        if (current_width + c.width().unwrap_or(1)) > allowed_width {
-            break;
+    let mut iter = console::AnsiCodeIterator::new(word);
+    let mut current_len = 0;
+
+    let mut escapes = Vec::new();
+    let mut head = String::with_capacity(word.len());
+    let mut last_txt = 0;
+    let mut escape_count = 0;
+    let mut tail = String::with_capacity(word.len());
+
+    for (val, is_esc) in iter.by_ref() {
+        if is_esc {
+            escapes.push(val);
+            if val == ANSI_RESET {
+                escapes.clear();
+            }
         }
 
-        // We can unwrap, as we just checked that a suitable character is next in line.
-        let c = char_iter.next().unwrap();
+        let len = match is_esc {
+            true => 0,
+            false => val.width(),
+        };
+
+        if current_len + len <= allowed_width {
+            head.push_str(val);
+            current_len += len;
 
-        // We default to 1 char, if the character length cannot be determined.
-        // The user has to live with this, if they decide to add control characters or some fancy
-        // stuff into their tables. This is considered undefined behavior and we try to handle this
-        // to the best of our capabilities.
-        let character_width = c.width().unwrap_or(1);
+            if !is_esc {
+                //allows popping unneeded escape codes later
+                last_txt = head.len();
+                escape_count = escapes.len();
+            }
+        } else {
+            assert!(!is_esc);
+            let mut char_iter = val.chars().peekable();
+            while let Some(c) = char_iter.peek() {
+                let character_width = c.width().unwrap_or(0);
+                if allowed_width < current_len + character_width {
+                    break;
+                }
 
-        current_width += character_width;
-        splitted.push(c);
+                current_len += character_width;
+                let c = char_iter.next().unwrap();
+                head.push(c);
+
+                // c is not escape code
+                last_txt = head.len();
+                escape_count = escapes.len();
+            }
+
+            head.truncate(last_txt); // cut off dangling escape codes since they should have no effect
+            if escape_count != 0 {
+                head.push_str(ANSI_RESET);
+            }
+
+            for esc in escapes {
+                tail.push_str(esc);
+            }
+            let remaining: String = char_iter.collect();
+            tail.push_str(&remaining);
+            break;
+        }
     }
 
-    // Collect the remaining characters.
-    let remaining = char_iter.collect();
-    (splitted, remaining)
+    iter.for_each(|s| tail.push_str(s.0));
+    (head, tail)
+}
+
+/// Fixes ansi escape codes in a split string
+/// 1. Adds reset code to the end of each substring if needed.
+/// 2. Keeps track of previous substring's escape codes and inserts them in later substrings to continue style
+pub fn fix_style_in_split_str(words: &mut [String]) {
+    let mut escapes: Vec<String> = Vec::new();
+
+    for word in words {
+        if !escapes.is_empty() {
+            word.insert_str(0, escapes.join("").as_str())
+        }
+
+        let iter = console::AnsiCodeIterator::new(word)
+            .filter(|(_, is_esc)| *is_esc)
+            .map(|v| v.0);
+        for esc in iter {
+            if esc == ANSI_RESET {
+                escapes.clear()
+            } else {
+                escapes.push(esc.to_string())
+            }
+        }
+
+        if !escapes.is_empty() {
+            word.push_str(ANSI_RESET);
+        }
+    }
 }

src/utils/mod.rs

@@ -7,6 +7,7 @@ use crate::{Column, Table};
 use arrangement::arrange_content;
 use formatting::borders::draw_borders;
 use formatting::content_format::format_content;
+pub use formatting::content_split::fix_style_in_split_str;
 
 /// This struct is ONLY used when table.to_string() is called.
 /// It's purpose is to store intermediate results, information on how to