From 6e0e10bad1cc3a8dee8c1202692ddf7b7f488fdb Mon Sep 17 00:00:00 2001
From: Nasonov V <qxuken@gmail.com>
Date: Wed, 2 Aug 2023 14:06:50 +0300
Subject: [PATCH] Added dev mode and extension filter for RustEmbed

---
 src/registry.rs | 56 +++++++++++++++++++++++++++++++++++++++++++------
 src/sources.rs  | 21 ++++++++++++++++++-
 tests/embed.rs  | 47 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 117 insertions(+), 7 deletions(-)

diff --git a/src/registry.rs b/src/registry.rs
index d1a69de41..02301e6fe 100644
--- a/src/registry.rs
+++ b/src/registry.rs
@@ -29,6 +29,8 @@ use rhai::Engine;
 #[cfg(feature = "script_helper")]
 use crate::helpers::scripting::ScriptHelper;
 
+#[cfg(feature = "rust-embed")]
+use crate::sources::LazySource;
 #[cfg(feature = "rust-embed")]
 use rust_embed::RustEmbed;
 
@@ -342,12 +344,14 @@ impl<'reg> Registry<'reg> {
 
     /// Register templates using a
     /// [RustEmbed](https://github.com/pyros2097/rust-embed) type
+    /// Calls register_embed_templates_with_extension with empty extension.
     ///
     /// File names from embed struct are used as template name.
     ///
     /// ```skip
     /// #[derive(RustEmbed)]
     /// #[folder = "templates"]
+    /// #[include = "*.hbs"]
     /// struct Assets;
     ///
     /// let mut hbs = Handlebars::new();
@@ -360,13 +364,53 @@ impl<'reg> Registry<'reg> {
     where
         E: RustEmbed,
     {
-        for item in E::iter() {
-            let file_name = item.as_ref();
-            if let Some(file) = E::get(file_name) {
-                let data = file.data;
+        self.register_embed_templates_with_extension::<E>("")
+    }
+
+    /// Register templates using a
+    /// [RustEmbed](https://github.com/pyros2097/rust-embed) type
+    /// * `tpl_extension`: the template file extension
+    ///
+    /// File names from embed struct are used as template name, but extension is stripped.
+    ///
+    /// When dev_mode enabled templates is reloaded
+    /// from embed struct everytime it's visied.
+    ///
+    /// ```skip
+    /// #[derive(RustEmbed)]
+    /// #[folder = "templates"]
+    /// struct Assets;
+    ///
+    /// let mut hbs = Handlebars::new();
+    /// hbs.register_embed_templates_with_extension::<Assets>(".hbs");
+    /// ```
+    ///
+    #[cfg(feature = "rust-embed")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "rust-embed")))]
+    pub fn register_embed_templates_with_extension<E>(
+        &mut self,
+        tpl_extension: &str,
+    ) -> Result<(), TemplateError>
+    where
+        E: RustEmbed,
+    {
+        for file_name in E::iter().filter(|x| x.ends_with(tpl_extension)) {
+            let tpl_name = file_name
+                .strip_suffix(tpl_extension)
+                .unwrap_or(&file_name)
+                .to_owned();
+            let source = LazySource::new(move || {
+                E::get(&file_name)
+                    .map(|file| file.data.to_vec())
+                    .and_then(|data| String::from_utf8(data).ok())
+            });
+            let tpl_content = source
+                .load()
+                .map_err(|e| (e, "Template load error".to_owned()))?;
+            self.register_template_string(&tpl_name, &tpl_content)?;
 
-                let tpl_content = String::from_utf8_lossy(data.as_ref());
-                self.register_template_string(file_name, tpl_content)?;
+            if self.dev_mode {
+                self.template_sources.insert(tpl_name, Arc::new(source));
             }
         }
         Ok(())
diff --git a/src/sources.rs b/src/sources.rs
index 8c8b2ba57..d0768e9a1 100644
--- a/src/sources.rs
+++ b/src/sources.rs
@@ -1,5 +1,5 @@
 use std::fs::File;
-use std::io::{BufReader, Error as IOError, Read};
+use std::io::{BufReader, Error as IOError, ErrorKind, Read};
 use std::path::PathBuf;
 
 pub(crate) trait Source {
@@ -32,3 +32,22 @@ impl Source for FileSource {
         Ok(buf)
     }
 }
+
+pub(crate) struct LazySource<F: Fn() -> Option<String>> {
+    loader: F,
+}
+
+impl<F: Fn() -> Option<String>> LazySource<F> {
+    pub(crate) fn new(loader: F) -> LazySource<F> {
+        LazySource { loader }
+    }
+}
+
+impl<F: Fn() -> Option<String>> Source for LazySource<F> {
+    type Item = String;
+    type Error = IOError;
+
+    fn load(&self) -> Result<Self::Item, Self::Error> {
+        (self.loader)().ok_or(IOError::new(ErrorKind::Other, "Source load error"))
+    }
+}
diff --git a/tests/embed.rs b/tests/embed.rs
index c803c7b19..8594752e2 100644
--- a/tests/embed.rs
+++ b/tests/embed.rs
@@ -25,3 +25,50 @@ fn test_embed() {
         "Hello, Andy"
     );
 }
+
+#[test]
+#[cfg(feature = "rust-embed")]
+fn test_embed_with_extension() {
+    use rust_embed::RustEmbed;
+
+    #[derive(RustEmbed)]
+    #[folder = "tests/templates/"]
+    struct Templates;
+
+    let mut hbs = handlebars::Handlebars::new();
+    hbs.register_embed_templates_with_extension::<Templates>(".hbs")
+        .unwrap();
+
+    assert_eq!(1, hbs.get_templates().len());
+
+    let data = json!({
+        "name": "Andy"
+    });
+
+    assert_eq!(hbs.render("hello", &data).unwrap().trim(), "Hello, Andy");
+}
+
+#[test]
+#[cfg(feature = "rust-embed")]
+fn test_embed_with_extension_and_tests_struct_root() {
+    use rust_embed::RustEmbed;
+
+    #[derive(RustEmbed)]
+    #[folder = "tests/"]
+    struct Templates;
+
+    let mut hbs = handlebars::Handlebars::new();
+    hbs.register_embed_templates_with_extension::<Templates>(".hbs")
+        .unwrap();
+
+    assert_eq!(1, hbs.get_templates().len());
+
+    let data = json!({
+        "name": "Andy"
+    });
+
+    assert_eq!(
+        hbs.render("templates/hello", &data).unwrap().trim(),
+        "Hello, Andy"
+    );
+}