feat: add +image to code blocks to consume their output as an image

Author: mfontanini

src/code/execute.rs

@@ -10,7 +10,7 @@ use std::{
     collections::{BTreeMap, HashMap},
     fmt::{self, Debug},
     fs::File,
-    io::{self, BufRead, BufReader, Write},
+    io::{self, BufRead, BufReader, Read, Write},
     path::{Path, PathBuf},
     process::{self, Child, Stdio},
     sync::{Arc, Mutex},
@@ -59,12 +59,17 @@ impl SnippetExecutor {
         let config = self.language_config(snippet)?;
         let script_dir = Self::write_snippet(snippet, config)?;
         let state: Arc<Mutex<ExecutionState>> = Default::default();
+        let output_type = match snippet.attributes.image {
+            true => OutputType::Binary,
+            false => OutputType::Lines,
+        };
         let reader_handle = CommandsRunner::spawn(
             state.clone(),
             script_dir,
             config.commands.clone(),
             config.environment.clone(),
             self.cwd.to_path_buf(),
+            output_type,
         );
         let handle = ExecutionHandle { state, reader_handle };
         Ok(handle)
@@ -183,15 +188,16 @@ impl CommandsRunner {
         commands: Vec<Vec<String>>,
         env: HashMap<String, String>,
         cwd: PathBuf,
+        output_type: OutputType,
     ) -> thread::JoinHandle<()> {
         let reader = Self { state, script_directory };
-        thread::spawn(|| reader.run(commands, env, cwd))
+        thread::spawn(move || reader.run(commands, env, cwd, output_type))
     }
 
-    fn run(self, commands: Vec<Vec<String>>, env: HashMap<String, String>, cwd: PathBuf) {
+    fn run(self, commands: Vec<Vec<String>>, env: HashMap<String, String>, cwd: PathBuf, output_type: OutputType) {
         let mut last_result = true;
         for command in commands {
-            last_result = self.run_command(command, &env, &cwd);
+            last_result = self.run_command(command, &env, &cwd, output_type);
             if !last_result {
                 break;
             }
@@ -203,17 +209,23 @@ impl CommandsRunner {
         self.state.lock().unwrap().status = status;
     }
 
-    fn run_command(&self, command: Vec<String>, env: &HashMap<String, String>, cwd: &Path) -> bool {
+    fn run_command(
+        &self,
+        command: Vec<String>,
+        env: &HashMap<String, String>,
+        cwd: &Path,
+        output_type: OutputType,
+    ) -> bool {
         let (mut child, reader) = match self.launch_process(command, env, cwd) {
             Ok(inner) => inner,
             Err(e) => {
                 let mut state = self.state.lock().unwrap();
                 state.status = ProcessStatus::Failure;
-                state.output.push(e.to_string());
+                state.output.extend(e.to_string().into_bytes());
                 return false;
             }
         };
-        let _ = Self::process_output(self.state.clone(), reader);
+        let _ = Self::process_output(self.state.clone(), reader, output_type);
 
         match child.wait() {
             Ok(code) => code.success(),
@@ -246,24 +258,41 @@ impl CommandsRunner {
         Ok((child, reader))
     }
 
-    fn process_output(state: Arc<Mutex<ExecutionState>>, reader: os_pipe::PipeReader) -> io::Result<()> {
-        let reader = BufReader::new(reader);
-        for line in reader.lines() {
-            let mut line = line?;
-            if line.contains('\t') {
-                line = line.replace('\t', "    ");
+    fn process_output(
+        state: Arc<Mutex<ExecutionState>>,
+        mut reader: os_pipe::PipeReader,
+        output_type: OutputType,
+    ) -> io::Result<()> {
+        match output_type {
+            OutputType::Lines => {
+                let reader = BufReader::new(reader);
+                for line in reader.lines() {
+                    let mut state = state.lock().unwrap();
+                    state.output.extend(line?.into_bytes());
+                    state.output.push(b'\n');
+                }
+                Ok(())
+            }
+            OutputType::Binary => {
+                let mut buffer = Vec::new();
+                reader.read_to_end(&mut buffer)?;
+                state.lock().unwrap().output.extend(buffer);
+                Ok(())
             }
-            // TODO: consider not locking per line...
-            state.lock().unwrap().output.push(line);
         }
-        Ok(())
     }
 }
 
+#[derive(Clone, Copy)]
+enum OutputType {
+    Lines,
+    Binary,
+}
+
 /// The state of the execution of a process.
 #[derive(Clone, Default, Debug)]
 pub(crate) struct ExecutionState {
-    pub(crate) output: Vec<String>,
+    pub(crate) output: Vec<u8>,
     pub(crate) status: ProcessStatus,
 }
 
@@ -307,8 +336,8 @@ echo 'bye'"
             }
         };
 
-        let expected_lines = vec!["hello world", "bye"];
-        assert_eq!(state.output, expected_lines);
+        let expected = b"hello world\nbye\n";
+        assert_eq!(state.output, expected);
     }
 
     #[test]
@@ -343,8 +372,8 @@ echo 'hello world'
             }
         };
 
-        let expected_lines = vec!["This message redirects to stderr", "hello world"];
-        assert_eq!(state.output, expected_lines);
+        let expected = b"This message redirects to stderr\nhello world\n";
+        assert_eq!(state.output, expected);
     }
 
     #[test]
@@ -368,9 +397,8 @@ echo 'hello world'
             }
         };
 
-        let expected_lines =
-            vec!["this line was hidden", "this line was hidden and contains another prefix /// ", "hello world"];
-        assert_eq!(state.output, expected_lines);
+        let expected = b"this line was hidden\nthis line was hidden and contains another prefix /// \nhello world\n";
+        assert_eq!(state.output, expected);
     }
 
     #[test]

src/code/snippet.rs

@@ -235,6 +235,10 @@ impl SnippetParser {
                 SnippetAttribute::LineNumbers => attributes.line_numbers = true,
                 SnippetAttribute::Exec => attributes.execute = true,
                 SnippetAttribute::ExecReplace => attributes.execute_replace = true,
+                SnippetAttribute::Image => {
+                    attributes.execute_replace = true;
+                    attributes.image = true;
+                }
                 SnippetAttribute::Render => attributes.render = true,
                 SnippetAttribute::NoBackground => attributes.no_background = true,
                 SnippetAttribute::AcquireTerminal => attributes.acquire_terminal = true,
@@ -259,6 +263,7 @@ impl SnippetParser {
                     "line_numbers" => SnippetAttribute::LineNumbers,
                     "exec" => SnippetAttribute::Exec,
                     "exec_replace" => SnippetAttribute::ExecReplace,
+                    "image" => SnippetAttribute::Image,
                     "render" => SnippetAttribute::Render,
                     "no_background" => SnippetAttribute::NoBackground,
                     "acquire_terminal" => SnippetAttribute::AcquireTerminal,
@@ -373,6 +378,7 @@ enum SnippetAttribute {
     LineNumbers,
     Exec,
     ExecReplace,
+    Image,
     Render,
     HighlightedLines(Vec<HighlightGroup>),
     Width(Percent),
@@ -570,6 +576,10 @@ pub(crate) struct SnippetAttributes {
     /// of its execution.
     pub(crate) execute_replace: bool,
 
+    /// Whether the snippet should be executed and its output should be considered to be an image
+    /// and replaced with it.
+    pub(crate) image: bool,
+
     /// Whether a snippet is marked to be rendered.
     ///
     /// A rendered snippet is transformed during parsing, leading to some visual

src/presentation/builder.rs

@@ -36,7 +36,8 @@ use crate::{
     third_party::{ThirdPartyRender, ThirdPartyRenderError, ThirdPartyRenderRequest},
     ui::{
         execution::{
-            DisplaySeparator, RunAcquireTerminalSnippet, RunSnippetOperation, SnippetExecutionDisabledOperation,
+            DisplaySeparator, RunAcquireTerminalSnippet, RunImageSnippet, RunSnippetOperation,
+            SnippetExecutionDisabledOperation,
         },
         footer::{FooterContext, FooterGenerator},
         modals::{IndexBuilder, KeyBindingsModalBuilder},
@@ -975,18 +976,31 @@ impl<'a> PresentationBuilder<'a> {
         style
     }
 
-    fn push_code_execution(&mut self, code: Snippet, block_length: usize, mode: ExecutionMode) -> BuildResult {
-        if !self.code_executor.is_execution_supported(&code.language) {
-            return Err(BuildError::UnsupportedExecution(code.language));
+    fn push_code_execution(&mut self, snippet: Snippet, block_length: usize, mode: ExecutionMode) -> BuildResult {
+        if !self.code_executor.is_execution_supported(&snippet.language) {
+            return Err(BuildError::UnsupportedExecution(snippet.language));
         }
-        if code.attributes.acquire_terminal {
+        if snippet.attributes.image {
+            let operation = RunImageSnippet::new(
+                snippet,
+                self.code_executor.clone(),
+                self.image_registry.clone(),
+                self.theme.execution_output.status.clone(),
+            );
+            operation.start_render();
+
+            let operation = RenderOperation::RenderAsync(Rc::new(operation));
+            self.chunk_operations.push(operation);
+            return Ok(());
+        }
+        if snippet.attributes.acquire_terminal {
             let block_length = block_length as u16;
             let block_length = match self.theme.code.alignment.clone().unwrap_or_default() {
                 Alignment::Left { .. } | Alignment::Right { .. } => block_length,
                 Alignment::Center { minimum_size, .. } => block_length.max(minimum_size),
             };
             let operation = RunAcquireTerminalSnippet::new(
-                code,
+                snippet,
                 self.code_executor.clone(),
                 self.theme.execution_output.status.clone(),
                 block_length,
@@ -999,14 +1013,14 @@ impl<'a> PresentationBuilder<'a> {
             ExecutionMode::AlongSnippet => DisplaySeparator::On,
             ExecutionMode::ReplaceSnippet => DisplaySeparator::Off,
         };
-        let alignment = self.code_style(&code).alignment.unwrap_or_default();
+        let alignment = self.code_style(&snippet).alignment.unwrap_or_default();
         let default_colors = self.theme.default_style.colors;
         let mut execution_output_style = self.theme.execution_output.clone();
-        if code.attributes.no_background {
+        if snippet.attributes.no_background {
             execution_output_style.colors.background = None;
         }
         let operation = RunSnippetOperation::new(
-            code,
+            snippet,
             self.code_executor.clone(),
             default_colors,
             execution_output_style,

src/terminal/image/printer.rs

@@ -14,7 +14,7 @@ use crate::{
 use image::{DynamicImage, ImageError};
 use std::{
     borrow::Cow,
-    io,
+    fmt, io,
     path::{Path, PathBuf},
     sync::Arc,
 };
@@ -162,6 +162,18 @@ impl PrintImage for ImagePrinter {
 #[derive(Clone, Default)]
 pub(crate) struct ImageRegistry(pub Arc<ImagePrinter>);
 
+impl fmt::Debug for ImageRegistry {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        let inner = match self.0.as_ref() {
+            ImagePrinter::Kitty(_) => "Kitty",
+            ImagePrinter::Iterm(_) => "Iterm",
+            ImagePrinter::Ascii(_) => "Ascii",
+            ImagePrinter::Null => "Null",
+        };
+        write!(f, "ImageRegistry<{inner}>")
+    }
+}
+
 impl ImageRegistry {
     pub(crate) fn register_image(&self, image: DynamicImage) -> Result<Image, RegisterImageError> {
         let resource = self.0.register(image)?;