Merge pull request #2618 from Piturnah/feat/write-to-buffer

feat: Optionally write to a buffer instead of stdout

Author: sharkdp

examples/buffer.rs

@@ -0,0 +1,17 @@
+use bat::{assets::HighlightingAssets, config::Config, controller::Controller, Input};
+
+fn main() {
+    let mut buffer = String::new();
+    let config = Config {
+        colored_output: true,
+        ..Default::default()
+    };
+    let assets = HighlightingAssets::from_binary();
+    let controller = Controller::new(&config, &assets);
+    let input = Input::from_file(file!());
+    controller
+        .run(vec![input.into()], Some(&mut buffer))
+        .unwrap();
+
+    println!("{buffer}");
+}

src/bin/bat/main.rs

@@ -206,7 +206,7 @@ pub fn list_themes(cfg: &Config, config_dir: &Path, cache_dir: &Path) -> Result<
             )?;
             config.theme = theme.to_string();
             Controller::new(&config, &assets)
-                .run(vec![theme_preview_file()])
+                .run(vec![theme_preview_file()], None)
                 .ok();
             writeln!(stdout)?;
         }
@@ -230,7 +230,7 @@ pub fn list_themes(cfg: &Config, config_dir: &Path, cache_dir: &Path) -> Result<
 fn run_controller(inputs: Vec<Input>, config: &Config, cache_dir: &Path) -> Result<bool> {
     let assets = assets_from_cache_or_binary(config.use_custom_assets, cache_dir)?;
     let controller = Controller::new(config, &assets);
-    controller.run(inputs)
+    controller.run(inputs, None)
 }
 
 #[cfg(feature = "bugreport")]

src/controller.rs

@@ -12,7 +12,7 @@ use crate::line_range::{LineRanges, RangeCheckResult};
 use crate::output::OutputType;
 #[cfg(feature = "paging")]
 use crate::paging::PagingMode;
-use crate::printer::{InteractivePrinter, Printer, SimplePrinter};
+use crate::printer::{InteractivePrinter, OutputHandle, Printer, SimplePrinter};
 
 use clircle::{Clircle, Identifier};
 
@@ -26,13 +26,18 @@ impl<'b> Controller<'b> {
         Controller { config, assets }
     }
 
-    pub fn run(&self, inputs: Vec<Input>) -> Result<bool> {
-        self.run_with_error_handler(inputs, default_error_handler)
+    pub fn run(
+        &self,
+        inputs: Vec<Input>,
+        output_buffer: Option<&mut dyn std::fmt::Write>,
+    ) -> Result<bool> {
+        self.run_with_error_handler(inputs, output_buffer, default_error_handler)
     }
 
     pub fn run_with_error_handler(
         &self,
         inputs: Vec<Input>,
+        output_buffer: Option<&mut dyn std::fmt::Write>,
         handle_error: impl Fn(&Error, &mut dyn Write),
     ) -> Result<bool> {
         let mut output_type;
@@ -74,24 +79,34 @@ impl<'b> Controller<'b> {
             clircle::Identifier::stdout()
         };
 
-        let writer = output_type.handle()?;
+        let mut writer = match output_buffer {
+            Some(buf) => OutputHandle::FmtWrite(buf),
+            None => OutputHandle::IoWrite(output_type.handle()?),
+        };
         let mut no_errors: bool = true;
         let stderr = io::stderr();
 
         for (index, input) in inputs.into_iter().enumerate() {
             let identifier = stdout_identifier.as_ref();
             let is_first = index == 0;
             let result = if input.is_stdin() {
-                self.print_input(input, writer, io::stdin().lock(), identifier, is_first)
+                self.print_input(input, &mut writer, io::stdin().lock(), identifier, is_first)
             } else {
                 // Use dummy stdin since stdin is actually not used (#1902)
-                self.print_input(input, writer, io::empty(), identifier, is_first)
+                self.print_input(input, &mut writer, io::empty(), identifier, is_first)
             };
             if let Err(error) = result {
-                if attached_to_pager {
-                    handle_error(&error, writer);
-                } else {
-                    handle_error(&error, &mut stderr.lock());
+                match writer {
+                    // It doesn't make much sense to send errors straight to stderr if the user
+                    // provided their own buffer, so we just return it.
+                    OutputHandle::FmtWrite(_) => return Err(error),
+                    OutputHandle::IoWrite(ref mut writer) => {
+                        if attached_to_pager {
+                            handle_error(&error, writer);
+                        } else {
+                            handle_error(&error, &mut stderr.lock());
+                        }
+                    }
                 }
                 no_errors = false;
             }
@@ -103,7 +118,7 @@ impl<'b> Controller<'b> {
     fn print_input<R: BufRead>(
         &self,
         input: Input,
-        writer: &mut dyn Write,
+        writer: &mut OutputHandle,
         stdin: R,
         stdout_identifier: Option<&Identifier>,
         is_first: bool,
@@ -164,7 +179,7 @@ impl<'b> Controller<'b> {
     fn print_file(
         &self,
         printer: &mut dyn Printer,
-        writer: &mut dyn Write,
+        writer: &mut OutputHandle,
         input: &mut OpenedInput,
         add_header_padding: bool,
         #[cfg(feature = "git")] line_changes: &Option<LineChanges>,
@@ -202,7 +217,7 @@ impl<'b> Controller<'b> {
     fn print_file_ranges(
         &self,
         printer: &mut dyn Printer,
-        writer: &mut dyn Write,
+        writer: &mut OutputHandle,
         reader: &mut InputReader,
         line_ranges: &LineRanges,
     ) -> Result<()> {

src/error.rs

@@ -7,6 +7,8 @@ pub enum Error {
     #[error(transparent)]
     Io(#[from] ::std::io::Error),
     #[error(transparent)]
+    Fmt(#[from] ::std::fmt::Error),
+    #[error(transparent)]
     SyntectError(#[from] ::syntect::Error),
     #[error(transparent)]
     SyntectLoadingError(#[from] ::syntect::LoadingError),

src/pretty_printer.rs

@@ -300,7 +300,7 @@ impl<'a> PrettyPrinter<'a> {
 
         // Run the controller
         let controller = Controller::new(&self.config, &self.assets);
-        controller.run(inputs.into_iter().map(|i| i.into()).collect())
+        controller.run(inputs.into_iter().map(|i| i.into()).collect(), None)
     }
 }