feat: process colored output from execution output

Cargo.toml

@@ -8,6 +8,7 @@ version = "0.8.0"
 edition = "2021"
 
 [dependencies]
+ansi-parser = "0.9"
 base64 = "0.22"
 bincode = "1.3"
 clap = { version = "4.4", features = ["derive", "string"] }

src/markdown/text.rs

@@ -207,7 +207,7 @@ impl<'a> WeightedTextRef<'a> {
         Self { text, accumulators, style: self.style }
     }
 
-    fn width(&self) -> usize {
+    pub(crate) fn width(&self) -> usize {
         let last_width = self.accumulators.last().map(|a| a.width).unwrap_or(0);
         let first_width = self.accumulators.first().map(|a| a.width).unwrap_or(0);
         last_width - first_width

src/processing/execution.rs

@@ -1,23 +1,25 @@
 use super::separator::{RenderSeparator, SeparatorWidth};
 use crate::{
     execute::{ExecutionHandle, ExecutionState, ProcessStatus, SnippetExecutor},
-    markdown::elements::{Snippet, Text, TextBlock},
+    markdown::{
+        elements::{Snippet, Text, TextBlock},
+        text::WeightedTextBlock,
+    },
     presentation::{AsRenderOperations, BlockLine, BlockLineText, RenderAsync, RenderAsyncState, RenderOperation},
-    render::properties::WindowSize,
+    render::{highlighting::AnsiSplitter, properties::WindowSize},
     style::{Colors, TextStyle},
     theme::{Alignment, ExecutionStatusBlockStyle},
     PresentationTheme,
 };
-use itertools::Itertools;
 use std::{cell::RefCell, mem, rc::Rc};
-use unicode_width::UnicodeWidthStr;
 
 #[derive(Debug)]
 struct RunSnippetOperationInner {
     handle: Option<ExecutionHandle>,
-    output_lines: Vec<String>,
+    output_lines: Vec<WeightedTextBlock>,
     state: RenderAsyncState,
     max_line_length: u16,
+    starting_style: TextStyle,
 }
 
 #[derive(Debug)]
@@ -54,6 +56,7 @@ impl RunSnippetOperation {
             output_lines: Vec::new(),
             state: RenderAsyncState::default(),
             max_line_length: 0,
+            starting_style: TextStyle::default(),
         };
         Self {
             code,
@@ -67,20 +70,10 @@ impl RunSnippetOperation {
             state_description: Text::new("running", TextStyle::default().colors(running_colors)).into(),
         }
     }
-
-    fn render_line(&self, line: String, block_length: u16) -> RenderOperation {
-        let line_len = line.width() as u16;
-        RenderOperation::RenderBlockLine(BlockLine {
-            text: BlockLineText::Preformatted(line),
-            unformatted_length: line_len,
-            block_length,
-            alignment: self.alignment.clone(),
-        })
-    }
 }
 
 impl AsRenderOperations for RunSnippetOperation {
-    fn as_render_operations(&self, dimensions: &WindowSize) -> Vec<RenderOperation> {
+    fn as_render_operations(&self, _dimensions: &WindowSize) -> Vec<RenderOperation> {
         let inner = self.inner.borrow();
         if matches!(inner.state, RenderAsyncState::NotStarted) {
             return Vec::new();
@@ -102,11 +95,13 @@ impl AsRenderOperations for RunSnippetOperation {
 
         let block_length = self.block_length.max(inner.max_line_length.saturating_add(1));
         for line in &inner.output_lines {
-            let chunks = line.chars().chunks(dimensions.columns as usize);
-            for chunk in &chunks {
-                operations.push(self.render_line(chunk.collect(), block_length));
-                operations.push(RenderOperation::RenderLineBreak);
-            }
+            operations.push(RenderOperation::RenderBlockLine(BlockLine {
+                text: BlockLineText::Weighted(line.clone()),
+                unformatted_length: line.width() as u16,
+                block_length,
+                alignment: self.alignment.clone(),
+            }));
+            operations.push(RenderOperation::RenderLineBreak);
         }
         operations.push(RenderOperation::SetColors(self.default_colors.clone()));
         operations
@@ -140,10 +135,12 @@ impl RenderAsync for RunSnippetOperation {
             drop(state);
 
             let mut max_line_length = 0;
+            let (new_lines, style) = AnsiSplitter::new(inner.starting_style.clone()).split_lines(&new_lines);
             for line in &new_lines {
                 let width = u16::try_from(line.width()).unwrap_or(u16::MAX);
                 max_line_length = max_line_length.max(width);
             }
+            inner.starting_style = style;
             if is_finished {
                 inner.handle.take();
                 inner.state = RenderAsyncState::JustFinishedRendering;
@@ -168,7 +165,7 @@ impl RenderAsync for RunSnippetOperation {
                 true
             }
             Err(e) => {
-                inner.output_lines = vec![e.to_string()];
+                inner.output_lines = vec![WeightedTextBlock::from(e.to_string())];
                 inner.state = RenderAsyncState::Rendered;
                 true
             }

src/render/engine.rs

@@ -220,18 +220,23 @@ where
         // Pad this code block with spaces so we get a nice little rectangle.
         let until_right_edge = max_line_length.saturating_sub(*unformatted_length);
         match text {
-            BlockLineText::Preformatted(text) => self.terminal.print_line(text)?,
-            BlockLineText::Weighted(text) => self.render_text(text, alignment)?,
+            BlockLineText::Preformatted(text) => {
+                self.terminal.print_line(text)?;
+                // If this line is longer than the screen, our cursor wrapped around so we need to update
+                // the terminal.
+                if *unformatted_length > max_line_length {
+                    let lines_wrapped = *unformatted_length / max_line_length;
+                    let new_row = self.terminal.cursor_row + lines_wrapped;
+                    self.terminal.sync_cursor_row(new_row)?;
+                }
+                self.terminal.print_line(&" ".repeat(until_right_edge as usize))?;
+            }
+            BlockLineText::Weighted(text) => {
+                let positioning = Positioning { max_line_length, start_column };
+                let text_drawer = TextDrawer::new_block(text, positioning, &self.colors)?;
+                text_drawer.draw(self.terminal)?;
+            }
         };
-        self.terminal.print_line(&" ".repeat(until_right_edge as usize))?;
-
-        // If this line is longer than the screen, our cursor wrapped around so we need to update
-        // the terminal.
-        if *unformatted_length > max_line_length {
-            let lines_wrapped = *unformatted_length / max_line_length;
-            let new_row = self.terminal.cursor_row + lines_wrapped;
-            self.terminal.sync_cursor_row(new_row)?;
-        }
 
         // Restore colors
         self.apply_colors()?;

src/render/highlighting.rs

@@ -1,4 +1,12 @@
-use crate::{markdown::elements::SnippetLanguage, theme::CodeBlockStyle};
+use crate::{
+    markdown::{
+        elements::{SnippetLanguage, Text, TextBlock},
+        text::WeightedTextBlock,
+    },
+    style::{Color, TextStyle},
+    theme::CodeBlockStyle,
+};
+use ansi_parser::{AnsiParser, AnsiSequence, Output};
 use crossterm::{
     style::{SetBackgroundColor, SetForegroundColor},
     QueueableCommand,
@@ -256,6 +264,97 @@ fn to_ansi_color(color: syntect::highlighting::Color) -> Option<crossterm::style
     }
 }
 
+pub(crate) struct AnsiSplitter {
+    lines: Vec<WeightedTextBlock>,
+    current_line: TextBlock,
+    current_style: TextStyle,
+}
+
+impl AnsiSplitter {
+    pub(crate) fn new(current_style: TextStyle) -> Self {
+        Self { lines: Default::default(), current_line: Default::default(), current_style }
+    }
+
+    pub(crate) fn split_lines(mut self, lines: &[String]) -> (Vec<WeightedTextBlock>, TextStyle) {
+        for line in lines {
+            for p in line.ansi_parse() {
+                match p {
+                    Output::TextBlock(text) => {
+                        self.current_line.0.push(Text::new(text, self.current_style.clone()));
+                    }
+                    Output::Escape(s) => self.handle_escape(&s),
+                }
+            }
+            let current_line = std::mem::take(&mut self.current_line);
+            self.lines.push(current_line.into());
+        }
+
+        (self.lines, self.current_style)
+    }
+
+    fn handle_escape(&mut self, s: &AnsiSequence) {
+        match s {
+            AnsiSequence::SetGraphicsMode(code) => {
+                let code = GraphicsCode(code);
+                code.update(&mut self.current_style);
+            }
+            AnsiSequence::EraseDisplay => {
+                self.lines.clear();
+                self.current_line.0.clear();
+            }
+            _ => (),
+        }
+    }
+}
+
+struct GraphicsCode<'a>(&'a [u8]);
+
+impl<'a> GraphicsCode<'a> {
+    fn update(&self, style: &mut TextStyle) {
+        // RGB mode
+        let codes = self.0;
+        if codes.starts_with(&[38, 2]) || codes.starts_with(&[48, 2]) {
+            if codes.len() == 5 {
+                let color = Color::new(codes[2], codes[3], codes[4]);
+                if codes[0] == 38 {
+                    style.colors.foreground = Some(color);
+                } else {
+                    style.colors.background = Some(color);
+                }
+            }
+            return;
+        }
+        for value in codes {
+            match value {
+                0 => *style = TextStyle::default(),
+                1 => *style = style.clone().bold(),
+                3 => *style = style.clone().italics(),
+                4 => *style = style.clone().underlined(),
+                9 => *style = style.clone().strikethrough(),
+                30 => style.colors.foreground = Some(Color::Black),
+                40 => style.colors.background = Some(Color::Black),
+                31 => style.colors.foreground = Some(Color::Red),
+                41 => style.colors.background = Some(Color::Red),
+                32 => style.colors.foreground = Some(Color::Green),
+                42 => style.colors.background = Some(Color::Green),
+                33 => style.colors.foreground = Some(Color::Yellow),
+                43 => style.colors.background = Some(Color::Yellow),
+                34 => style.colors.foreground = Some(Color::Blue),
+                44 => style.colors.background = Some(Color::Blue),
+                35 => style.colors.foreground = Some(Color::Magenta),
+                45 => style.colors.background = Some(Color::Magenta),
+                36 => style.colors.foreground = Some(Color::Cyan),
+                46 => style.colors.background = Some(Color::Cyan),
+                37 => style.colors.foreground = Some(Color::White),
+                47 => style.colors.background = Some(Color::White),
+                39 => style.colors.foreground = None,
+                49 => style.colors.background = None,
+                _ => (),
+            }
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use super::*;