From d79cb9d8f0b860ac66605860332118e63afccebb Mon Sep 17 00:00:00 2001
From: Ning Sun <sunng@about.me>
Date: Sat, 16 Sep 2017 13:34:33 +0800
Subject: [PATCH] Add 'Template::custom()' to customize templating engines.

---
 contrib/src/lib.rs                            |  6 +++
 contrib/src/templates/engine.rs               |  4 +-
 contrib/src/templates/mod.rs                  | 37 +++++++++++++++++--
 examples/handlebars_templates/src/main.rs     | 17 ++++++++-
 .../templates/index.html.hbs                  |  1 +
 5 files changed, 58 insertions(+), 7 deletions(-)

diff --git a/contrib/src/lib.rs b/contrib/src/lib.rs
index d538fee5..146f1982 100644
--- a/contrib/src/lib.rs
+++ b/contrib/src/lib.rs
@@ -48,6 +48,12 @@ extern crate serde;
 #[cfg_attr(feature = "json", macro_reexport(json_internal))]
 extern crate serde_json;
 
+#[cfg(feature = "handlebars_templates")]
+pub extern crate handlebars;
+
+#[cfg(feature = "tera_templates")]
+pub extern crate tera;
+
 #[cfg(feature = "json")]
 #[cfg_attr(feature = "json", macro_use)]
 #[doc(hidden)]
diff --git a/contrib/src/templates/engine.rs b/contrib/src/templates/engine.rs
index 0136c4b1..55f84a02 100644
--- a/contrib/src/templates/engine.rs
+++ b/contrib/src/templates/engine.rs
@@ -15,9 +15,9 @@ pub trait Engine: Send + Sync + 'static {
 
 pub struct Engines {
     #[cfg(feature = "tera_templates")]
-    tera: Tera,
+    pub tera: Tera,
     #[cfg(feature = "handlebars_templates")]
-    handlebars: Handlebars,
+    pub handlebars: Handlebars,
 }
 
 impl Engines {
diff --git a/contrib/src/templates/mod.rs b/contrib/src/templates/mod.rs
index a2756ac5..e8d4888e 100644
--- a/contrib/src/templates/mod.rs
+++ b/contrib/src/templates/mod.rs
@@ -158,9 +158,35 @@ impl Template {
     /// }
     /// ```
     pub fn fairing() -> impl Fairing {
-        AdHoc::on_attach(|rocket| {
-            let mut template_root = rocket.config()
-                .root_relative(DEFAULT_TEMPLATE_DIR);
+        Template::custom(|_| {})
+    }
+
+    /// Returns a fairing that intializes and maintains templating state.
+    ///
+    /// This method allows you to configure the template context via the
+    /// closure.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// extern crate rocket;
+    /// extern crate rocket_contrib;
+    ///
+    /// use rocket_contrib::Template;
+    ///
+    /// fn main() {
+    ///     rocket::ignite()
+    ///         // ...
+    ///         .attach(Template::custom(|engines| {
+    ///             // engines.handlebars.register_helper ...
+    ///         }))
+    ///         // ...
+    ///     # ;
+    /// }
+    /// ```
+    pub fn custom<F>(f: F) -> impl Fairing where F: Fn(&mut Engines) + Send + Sync + 'static {
+        AdHoc::on_attach(move |rocket| {
+            let mut template_root = rocket.config().root_relative(DEFAULT_TEMPLATE_DIR);
 
             match rocket.config().get_str("template_dir") {
                 Ok(dir) => template_root = rocket.config().root_relative(dir),
@@ -172,7 +198,10 @@ impl Template {
             };
 
             match Context::initialize(template_root) {
-                Some(ctxt) => Ok(rocket.manage(ctxt)),
+                Some(mut ctxt) => {
+                    f(&mut ctxt.engines);
+                    Ok(rocket.manage(ctxt))
+                }
                 None => Err(rocket)
             }
         })
diff --git a/examples/handlebars_templates/src/main.rs b/examples/handlebars_templates/src/main.rs
index 330353c9..7acbd4ae 100644
--- a/examples/handlebars_templates/src/main.rs
+++ b/examples/handlebars_templates/src/main.rs
@@ -10,6 +10,7 @@ extern crate rocket;
 use rocket::Request;
 use rocket::response::Redirect;
 use rocket_contrib::Template;
+use rocket_contrib::handlebars::{Helper, Handlebars, RenderContext, RenderError, JsonRender};
 
 #[derive(Serialize)]
 struct TemplateContext {
@@ -42,7 +43,21 @@ fn not_found(req: &Request) -> Template {
 fn rocket() -> rocket::Rocket {
     rocket::ignite()
         .mount("/", routes![index, get])
-        .attach(Template::fairing())
+        .attach(Template::custom(|engines| {
+            engines.handlebars.register_helper(
+                "echo", Box::new(|h: &Helper,
+                                  _: &Handlebars,
+                                  rc: &mut RenderContext| -> Result<(), RenderError> {
+                                      if let Some(p0) = h.param(0) {
+                                          rc.writer.write(p0.value()
+                                                          .render()
+                                                          .into_bytes()
+                                                          .as_ref())?;
+                                      }
+                                      Ok(())
+                                  }));
+
+        }))
         .catch(catchers![not_found])
 }
 
diff --git a/examples/handlebars_templates/templates/index.html.hbs b/examples/handlebars_templates/templates/index.html.hbs
index 62eb85b7..90ebf069 100644
--- a/examples/handlebars_templates/templates/index.html.hbs
+++ b/examples/handlebars_templates/templates/index.html.hbs
@@ -14,5 +14,6 @@
     </ul>
 
     <p>Try going to <a href="/hello/YourName">/hello/YourName</a></p>
+    <p>You can define custom helper like <b>{{echo "this"}}</b>.</p>
   </body>
 </html>