Add optional support for Rocket

1 files changed, 24 insertions, 0 deletions

diff --git a/askama_derive/src/generator.rs b/askama_derive/src/generator.rs
index 3595a84..1d8e8f2 100644
--- a/askama_derive/src/generator.rs
+++ b/askama_derive/src/generator.rs
@@ -56,6 +56,9 @@ pub fn generate(ast: &syn::DeriveInput, path: &Path, mut nodes: Vec<Node>) -> St
     if cfg!(feature = "iron") {
         gen.impl_modifier_response(ast);
     }
+    if cfg!(feature = "rocket") {
+        gen.impl_responder(ast, path);
+    }
     gen.result()
 }
 
@@ -645,6 +648,27 @@ impl<'a> Generator<'a> {
         self.writeln("}");
     }
 
+    // Implement Rocket's `Responder`.
+    fn impl_responder(&mut self, ast: &syn::DeriveInput, path: &Path) {
+        self.write_header(ast, "::askama::rocket::Responder<'r>", &vec!["'r"]);
+        self.writeln("fn respond_to(self, _: &::askama::rocket::Request) \
+                      -> ::std::result::Result<\
+                         ::askama::rocket::Response<'r>, ::askama::rocket::Status> {");
+        let ext = match path.extension() {
+            Some(s) => s.to_str().unwrap(),
+            None => "txt",
+        };
+        self.writeln("::askama::rocket::Response::build()");
+        self.indent();
+        self.writeln(&format!(".header(::askama::rocket::ContentType::from_extension({:?})\
+                               .unwrap())", ext));
+        self.writeln(".sized_body(::std::io::Cursor::new(self.render().unwrap()))");
+        self.writeln(".ok()");
+        self.dedent();
+        self.writeln("}");
+        self.writeln("}");
+    }
+
     fn result(self) -> String {
         self.buf
     }