Move script formatting to Script type definition

This patch moves the script formatting to the definition of the Script
type, where it belongs to.

This is not an optimal implementation code-cleanness wise, because we
could always improve the interface to be even more typesafe and
convenient to use, but I guess this is nice for now.

Signed-off-by: Matthias Beyer <mail@beyermatthias.de>

3 files changed, 97 insertions, 50 deletions

diff --git a/src/commands/db.rs b/src/commands/db.rs
index eb65425..fa52d8a 100644
--- a/src/commands/db.rs
+++ b/src/commands/db.rs
@@ -20,6 +20,7 @@ use crate::config::Configuration;
 use crate::db::DbConnectionConfig;
 use crate::db::models;
 use crate::log::LogItem;
+use crate::package::Script;
 use crate::schema;
 
 pub fn db(db_connection_config: DbConnectionConfig, config: &Configuration, matches: &ArgMatches) -> Result<()> {
@@ -387,7 +388,11 @@ fn job(conn_cfg: DbConnectionConfig, config: &Configuration, matches: &ArgMatche
             log_len         = format!("{:<4}", data.0.log_text.lines().count()).cyan(),
             envs            = env_vars.unwrap_or_else(|| String::from("<env vars hidden>")),
             script_text     = if show_script {
-                crate::ui::script_to_printable(&data.0.script_text, script_highlight, configured_theme.as_ref(), script_line_numbers)?
+                let theme = configured_theme
+                    .as_ref()
+                    .ok_or_else(|| anyhow!("Highlighting for script enabled, but no theme configured"))?;
+                let script = Script::from(data.0.script_text);
+                crate::ui::script_to_printable(&script, script_highlight, theme, script_line_numbers)?
             } else {
                 String::from("<script hidden>")
             },
diff --git a/src/package/script.rs b/src/package/script.rs
index 1c92040..805d17f 100644
--- a/src/package/script.rs
+++ b/src/package/script.rs
@@ -1,8 +1,13 @@
+use anyhow::anyhow;
 use anyhow::Error;
 use anyhow::Result;
 use handlebars::{Handlebars, HelperDef, RenderContext, Helper, Context, JsonRender, HelperResult, Output, RenderError};
 use serde::Deserialize;
 use serde::Serialize;
+use syntect::easy::HighlightLines;
+use syntect::highlighting::{ThemeSet, Style};
+use syntect::parsing::SyntaxSet;
+use syntect::util::{as_24_bit_terminal_escaped, LinesWithEndings};
 
 use crate::package::Package;
 use crate::phase::Phase;
@@ -12,9 +17,78 @@ use crate::phase::PhaseName;
 #[serde(transparent)]
 pub struct Script(String);
 
+impl From<String> for Script {
+    fn from(s: String) -> Script {
+        Script(s)
+    }
+}
+
+impl std::fmt::Display for Script {
+    fn fmt(&self, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(fmt, "{}", self.0)
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct Shebang(String);
 
+impl Script {
+    pub fn highlighted<'a>(&'a self, script_theme: &'a str) -> HighlightedScript<'a> {
+        HighlightedScript::new(self, script_theme)
+    }
+
+    pub fn lines_numbered(&self) -> impl Iterator<Item = (usize, &str)> {
+        self.0.lines().enumerate()
+    }
+}
+
+#[derive(Debug)]
+pub struct HighlightedScript<'a> {
+    script: &'a Script,
+    script_theme: &'a str,
+
+    ps: SyntaxSet,
+    ts: ThemeSet,
+}
+
+impl<'a> HighlightedScript<'a> {
+    fn new(script: &'a Script, script_theme: &'a str) -> Self {
+        HighlightedScript {
+            script,
+            script_theme,
+
+            ps: SyntaxSet::load_defaults_newlines(),
+            ts: ThemeSet::load_defaults(),
+        }
+    }
+
+    pub fn lines(&'a self) -> Result<impl Iterator<Item = String> + 'a> {
+        let syntax = self.ps
+            .find_syntax_by_first_line(&self.script.0)
+            .ok_or_else(|| anyhow!("Failed to load syntax for highlighting script"))?;
+
+        let theme = self.ts
+            .themes
+            .get(self.script_theme)
+            .ok_or_else(|| anyhow!("Theme not available: {}", self.script_theme))?;
+
+        let mut h = HighlightLines::new(syntax, &theme);
+
+        Ok({
+            LinesWithEndings::from(&self.script.0)
+                .map(move |line| {
+                    let ranges: Vec<(Style, &str)> = h.highlight(line, &self.ps);
+                    as_24_bit_terminal_escaped(&ranges[..], true)
+                })
+        })
+    }
+
+
+    pub fn lines_numbered(&'a self) -> Result<impl Iterator<Item = (usize, String)> + 'a> {
+        self.lines().map(|iter| iter.enumerate())
+    }
+}
+
 impl From<String> for Shebang {
     fn from(s: String) -> Self {
         Shebang(s)
diff --git a/src/ui.rs b/src/ui.rs
index 315d097..c3be7f6 100644
--- a/src/ui.rs
+++ b/src/ui.rs
@@ -10,12 +10,9 @@ use anyhow::anyhow;
 use handlebars::Handlebars;
 use itertools::Itertools;
 use log::error;
-use syntect::easy::HighlightLines;
-use syntect::highlighting::{ThemeSet, Style};
-use syntect::parsing::SyntaxSet;
-use syntect::util::{as_24_bit_terminal_escaped, LinesWithEndings};
 
 use crate::package::Package;
+use crate::package::Script;
 use crate::package::ScriptBuilder;
 use crate::package::Shebang;
 use crate::config::Configuration;
@@ -95,9 +92,12 @@ fn print_package(out: &mut dyn Write,
     let script = ScriptBuilder::new(&Shebang::from(config.shebang().clone()))
         .build(package, config.available_phases(), *config.strict_script_interpolation())?;
 
-    let script = crate::ui::script_to_printable(script.as_ref(),
+    let script = crate::ui::script_to_printable(&script,
         flags.script_highlighting,
-        config.script_highlight_theme().as_ref(),
+        config
+            .script_highlight_theme()
+            .as_ref()
+            .ok_or_else(|| anyhow!("Highlighting for script enabled, but no theme configured"))?,
         flags.script_line_numbers)?;
 
     let mut data = BTreeMap::new();
@@ -123,56 +123,24 @@ fn print_package(out: &mut dyn Write,
         .and_then(|r| writeln!(out, "{}", r).map_err(Error::from))
 }
 
-pub fn script_to_printable(script: &str,
-                          highlight: bool,
-                          highlight_theme: Option<&String>,
-                          line_numbers: bool)
+pub fn script_to_printable(script: &Script, highlight: bool, highlight_theme: &str, line_numbers: bool)
     -> Result<String>
 {
-    if highlight {
-        if let Some(configured_theme) = highlight_theme {
-            // Load these once at the start of your program
-            let ps = SyntaxSet::load_defaults_newlines();
-            let ts = ThemeSet::load_defaults();
-
-            let syntax = ps
-                .find_syntax_by_first_line(script)
-                .ok_or_else(|| anyhow!("Failed to load syntax for highlighting script"))?;
-
-            let theme = ts
-                .themes
-                .get(configured_theme)
-                .ok_or_else(|| anyhow!("Theme not available: {}", configured_theme))?;
-
-            let mut h = HighlightLines::new(syntax, &theme);
-
-            let output = LinesWithEndings::from(script)
-                .enumerate()
-                .map(|(i, line)| {
-                    let ranges: Vec<(Style, &str)> = h.highlight(line, &ps);
-                    if line_numbers {
-                        format!("{:>4} | {}", i, as_24_bit_terminal_escaped(&ranges[..], true))
-                    } else {
-                        as_24_bit_terminal_escaped(&ranges[..], true)
-                    }
-                })
-                .join("");
-
-            Ok(output)
+    let script = if highlight {
+        let script = script.highlighted(highlight_theme);
+        if line_numbers {
+            script.lines_numbered()?.map(|(i, s)| format!("{:>4} | {}", i, s)).join("")
         } else {
-            Err(anyhow!("Highlighting for script enabled, but no theme configured"))
+            script.lines()?.join("")
         }
     } else {
         if line_numbers {
-            Ok({
-                script.lines()
-                    .enumerate()
-                    .map(|(i, s)| format!("{:>4} | {}", i, s))
-                    .join("\n")
-            })
+            script.lines_numbered().map(|(i, s)| format!("{:>4} | {}", i, s)).join("")
         } else {
-            Ok(script.to_string())
+            script.to_string()
         }
-    }
+    };
+
+    Ok(script)
 }