<commit_before>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::fmt::{self, Debug, Formatter};

use serde::Serialize;

use regex::{Captures, Regex};

use template::Template;
use render::{RenderContext, Renderable};
use context::Context;
use helpers::{self, HelperDef};
use directives::{self, DirectiveDef};
use support::str::StringWriter;
use error::{RenderError, TemplateError, TemplateFileError, TemplateRenderError};

lazy_static!{
    static ref DEFAULT_REPLACE: Regex = Regex::new(">|<|\"|&").unwrap();
}

/// This type represents an *escape fn*, that is a function who's purpose it is
/// to escape potentially problematic characters in a string.
///
/// An *escape fn* is represented as a `Box` to avoid unnecessary type
/// parameters (and because traits cannot be aliased using `type`).
pub type EscapeFn = Box<Fn(&str) -> String + Send + Sync>;

/// The default *escape fn* replaces the characters `&"<>`
/// with the equivalent html / xml entities.
pub fn html_escape(data: &str) -> String {
    DEFAULT_REPLACE
        .replace_all(data, |cap: &Captures| {
            match cap.get(0).map(|m| m.as_str()) {
                Some("<") => "<",
                Some(">") => ">",
                Some("\"") => """,
                Some("&") => "&",
                _ => unreachable!(),
            }.to_owned()
        })
        .into_owned()
}

/// `EscapeFn` that do not change any thing. Useful when using in a non-html
/// environment.
pub fn no_escape(data: &str) -> String {
    data.to_owned()
}

/// The single entry point of your Handlebars templates
///
/// It maintains compiled templates and registered helpers.
pub struct Registry {
    templates: HashMap<String, Template>,
    helpers: HashMap<String, Box<HelperDef + 'static>>,
    directives: HashMap<String, Box<DirectiveDef + 'static>>,
    escape_fn: EscapeFn,
    source_map: bool,
    strict_mode: bool,
}

impl Debug for Registry {
    fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
        f.debug_struct("Handlebars")
            .field("templates", &self.templates)
            .field("helpers", &self.helpers.keys())
            .field("directives", &self.directives.keys())
            .field("source_map", &self.source_map)
            .finish()
    }
}

impl Registry {
    pub fn new() -> Registry {
        let r = Registry {
            templates: HashMap::new(),
            helpers: HashMap::new(),
            directives: HashMap::new(),
            escape_fn: Box::new(html_escape),
            source_map: true,
            strict_mode: false,
        };

        r.setup_builtins()
    }

    fn setup_builtins(mut self) -> Registry {
        self.register_helper("if", Box::new(helpers::IF_HELPER));
        self.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        self.register_helper("each", Box::new(helpers::EACH_HELPER));
        self.register_helper("with", Box::new(helpers::WITH_HELPER));
        self.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        self.register_helper("raw", Box::new(helpers::RAW_HELPER));
        self.register_helper("log", Box::new(helpers::LOG_HELPER));

        self.register_decorator("inline", Box::new(directives::INLINE_DIRECTIVE));
        self
    }

    /// Enable handlebars template source map
    ///
    /// Source map provides line/col reporting on error. It uses slightly
    /// more memory to maintain the data.
    ///
    /// Default is true.
    pub fn source_map_enabled(&mut self, enable: bool) {
        self.source_map = enable;
    }

    /// Enable handlebars strict mode
    ///
    /// By default, handlebars renders empty string for value that
    /// undefined or never exists. Since rust is a static type
    /// language, we offer strict mode in handlebars-rust.  In strict
    /// mode, if you were access a value that doesn't exist, a
    /// `RenderError` will be raised.
    pub fn set_strict_mode(&mut self, enable: bool) {
        self.strict_mode = enable;
    }

    /// Return strict mode state, default is false.
    ///
    /// By default, handlebars renders empty string for value that
    /// undefined or never exists. Since rust is a static type
    /// language, we offer strict mode in handlebars-rust.  In strict
    /// mode, if you were access a value that doesn't exist, a
    /// `RenderError` will be raised.
    pub fn strict_mode(&self) -> bool {
        self.strict_mode
    }

    /// Register a template string
    ///
    /// Returns `TemplateError` if there is syntax error on parsing template.
    pub fn register_template_string<S>(
        &mut self,
        name: &str,
        tpl_str: S,
    ) -> Result<(), TemplateError>
    where
        S: AsRef<str>,
    {
        try!(
            Template::compile_with_name(tpl_str, name.to_owned(), self.source_map)
                .and_then(|t| Ok(self.templates.insert(name.to_string(), t)))
        );
        Ok(())
    }

    /// Register a partial string
    ///
    /// A named partial will be added to the registry. It will overwrite template with
    /// same name. Currently registered partial is just identical to template.
    pub fn register_partial<S>(&mut self, name: &str, partial_str: S) -> Result<(), TemplateError>
    where
        S: AsRef<str>,
    {
        self.register_template_string(name, partial_str)
    }

    /// Register a template from a path
    pub fn register_template_file
(
        &mut self,
        name: &str,
        tpl_path: P,
    ) -> Result<(), TemplateFileError>
    where
        P: AsRef<Path>,
    {
        let mut file =
            try!(File::open(tpl_path).map_err(|e| TemplateFileError::IOError(e, name.to_owned())));
        self.register_template_source(name, &mut file)
    }

    /// Register a template from `std::io::Read` source
    pub fn register_template_source(
        &mut self,
        name: &str,
        tpl_source: &mut Read,
    ) -> Result<(), TemplateFileError> {
        let mut buf = String::new();
        try!(
            tpl_source
                .read_to_string(&mut buf)
                .map_err(|e| TemplateFileError::IOError(e, name.to_owned()))
        );
        try!(self.register_template_string(name, buf));
        Ok(())
    }

    /// remove a template from the registry
    pub fn unregister_template(&mut self, name: &str) {
        self.templates.remove(name);
    }

    /// register a helper
    pub fn register_helper(
        &mut self,
        name: &str,
        def: Box<HelperDef + 'static>,
    ) -> Option<Box<HelperDef + 'static>> {
        self.helpers.insert(name.to_string(), def)
    }

    /// register a decorator
    pub fn register_decorator(
        &mut self,
        name: &str,
        def: Box<DirectiveDef + 'static>,
    ) -> Option<Box<DirectiveDef + 'static>> {
        self.directives.insert(name.to_string(), def)
    }

    /// Register a new *escape fn* to be used from now on by this registry.
    pub fn register_escape_fn<F: 'static + Fn(&str) -> String + Send + Sync>(
        &mut self,
        escape_fn: F,
    ) {
        self.escape_fn = Box::new(escape_fn);
    }

    /// Restore the default *escape fn*.
    pub fn unregister_escape_fn(&mut self) {
        self.escape_fn = Box::new(html_escape);
    }

    /// Get a reference to the current *escape fn*.
    pub fn get_escape_fn(&self) -> &Fn(&str) -> String {
        &*self.escape_fn
    }

    /// Return a registered template,
    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    /// Return a registered helper
    pub fn get_helper(&self, name: &str) -> Option<&Box<HelperDef + 'static>> {
        self.helpers.get(name)
    }

    /// Return a registered directive, aka decorator
    pub fn get_decorator(&self, name: &str) -> Option<&Box<DirectiveDef + 'static>> {
        self.directives.get(name)
    }

    /// Return all templates registered
    pub fn get_templates(&self) -> &HashMap<String, Template> {
        &self.templates
    }

    /// Unregister all templates
    pub fn clear_templates(&mut self) {
        self.templates.clear();
    }

    /// Render a registered template with some data into a string
    ///
    /// * `name` is the template name you registred previously
    /// * `ctx` is the data that implements `serde::Serialize`
    ///
    /// Returns rendered string or an struct with error information
    pub fn render<T>(&self, name: &str, data: &T) -> Result<String, RenderError>
    where
        T: Serialize,
    {
        let mut writer = StringWriter::new();
        {
            try!(self.render_to_write(name, data, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// Render a registered template and write some data to the `std::io::Write`
    pub fn render_to_write<T>(
        &self,
        name: &str,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), RenderError>
    where
        T: Serialize,
    {
        self.get_template(&name.to_string())
            .ok_or(RenderError::new(format!("Template not found: {}", name)))
            .and_then(|t| {
                let ctx = try!(Context::wraps(data));
                let mut local_helpers = HashMap::new();
                let mut render_context = RenderContext::new(ctx, &mut local_helpers, writer);
                render_context.root_template = t.name.clone();
                t.render(self, &mut render_context)
            })
    }

    /// render a template string using current registry without register it
    pub fn render_template<T>(
        &self,
        template_string: &str,
        data: &T,
    ) -> Result<String, TemplateRenderError>
    where
        T: Serialize,
    {
        let mut writer = StringWriter::new();
        {
            try!(self.render_template_to_write(template_string, data, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// render a template string using current registry without register it
    pub fn render_template_to_write<T>(
        &self,
        template_string: &str,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        let tpl = try!(Template::compile2(template_string, self.source_map));
        let ctx = try!(Context::wraps(data));
        let mut local_helpers = HashMap::new();
        let mut render_context = RenderContext::new(ctx, &mut local_helpers, writer);
        tpl.render(self, &mut render_context)
            .map_err(TemplateRenderError::from)
    }

    /// render a template source using current registry without register it
    pub fn render_template_source_to_write<T>(
        &self,
        template_source: &mut Read,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        let mut tpl_str = String::new();
        try!(
            template_source
                .read_to_string(&mut tpl_str)
                .map_err(|e| TemplateRenderError::IOError(e, "Unamed template source".to_owned()))
        );
        self.render_template_to_write(&tpl_str, data, writer)
    }

    #[deprecated(since = "0.30.0", note = "Please use render_to_write instead.")]
    pub fn renderw<T>(&self, name: &str, data: &T, writer: &mut Write) -> Result<(), RenderError>
    where
        T: Serialize,
    {
        self.render_to_write(name, data, writer)
    }

    #[deprecated(since = "0.30.0", note = "Please use render_template instead.")]
    pub fn template_render<T>(
        &self,
        template_string: &str,
        data: &T,
    ) -> Result<String, TemplateRenderError>
    where
        T: Serialize,
    {
        self.render_template(template_string, data)
    }

    #[deprecated(since = "0.30.0", note = "Please use render_template_to_write instead.")]
    pub fn template_renderw<T>(
        &self,
        template_string: &str,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        self.render_template_to_write(template_string, data, writer)
    }

    #[deprecated(since = "0.30.0", note = "Please use render_template_source_to_write instead.")]
    pub fn template_renderw2<T>(
        &self,
        template_source: &mut Read,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        self.render_template_source_to_write(template_source, data, writer)
    }
}

#[cfg(test)]
mod test {
    use registry::Registry;
    use render::{Helper, RenderContext, Renderable};
    use helpers::HelperDef;
    use support::str::StringWriter;
    use error::RenderError;

    #[derive(Clone, Copy)]
    struct DummyHelper;

    impl HelperDef for DummyHelper {
        fn call(
            &self,
            h: &Helper,
            r: &Registry,
            rc: &mut RenderContext,
        ) -> Result<(), RenderError> {
            try!(h.template().unwrap().render(r, rc));
            Ok(())
        }
    }

    static DUMMY_HELPER: DummyHelper = DummyHelper;

    #[test]
    fn test_registry_operations() {
        let mut r = Registry::new();

        assert!(r.register_template_string("index", "

").is_ok());
        assert!(r.register_template_string("index2", "

").is_ok());

        assert_eq!(r.templates.len(), 2);

        r.unregister_template("index");
        assert_eq!(r.templates.len(), 1);

        r.clear_templates();
        assert_eq!(r.templates.len(), 0);

        r.register_helper("dummy", Box::new(DUMMY_HELPER));

        // built-in helpers plus 1
        assert_eq!(r.helpers.len(), 7 + 1);
    }

    #[test]
    fn test_render_to_write() {
        let mut r = Registry::new();

        assert!(r.register_template_string("index", "

").is_ok());

        let mut sw = StringWriter::new();
        {
            r.render_to_write("index", &(), &mut sw).ok().unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());
    }

    #[test]
    fn test_escape_fn() {
        let mut r = Registry::new();

        let input = String::from("\"<>&");

        r.register_template_string("test", String::from("{{this}}"))
            .unwrap();

        assert_eq!(""<>&", r.render("test", &input).unwrap());

        r.register_escape_fn(|s| s.into());

        assert_eq!("\"<>&", r.render("test", &input).unwrap());

        r.unregister_escape_fn();

        assert_eq!(""<>&", r.render("test", &input).unwrap());
    }

    #[test]
    fn test_escape() {
        let r = Registry::new();
        let data = json!({
            "hello": "world"
        });

        assert_eq!(
            "{{hello}}",
            r.render_template(r"\\{{hello}}", &data).unwrap()
        );
    }
}
<commit_msg>(test) added test for strict mode<commit_after>use std::collections::HashMap;
use std::io::prelude::*;
use std::fs::File;
use std::path::Path;
use std::fmt::{self, Debug, Formatter};

use serde::Serialize;

use regex::{Captures, Regex};

use template::Template;
use render::{RenderContext, Renderable};
use context::Context;
use helpers::{self, HelperDef};
use directives::{self, DirectiveDef};
use support::str::StringWriter;
use error::{RenderError, TemplateError, TemplateFileError, TemplateRenderError};

lazy_static!{
    static ref DEFAULT_REPLACE: Regex = Regex::new(">|<|\"|&").unwrap();
}

/// This type represents an *escape fn*, that is a function who's purpose it is
/// to escape potentially problematic characters in a string.
///
/// An *escape fn* is represented as a `Box` to avoid unnecessary type
/// parameters (and because traits cannot be aliased using `type`).
pub type EscapeFn = Box<Fn(&str) -> String + Send + Sync>;

/// The default *escape fn* replaces the characters `&"<>`
/// with the equivalent html / xml entities.
pub fn html_escape(data: &str) -> String {
    DEFAULT_REPLACE
        .replace_all(data, |cap: &Captures| {
            match cap.get(0).map(|m| m.as_str()) {
                Some("<") => "<",
                Some(">") => ">",
                Some("\"") => """,
                Some("&") => "&",
                _ => unreachable!(),
            }.to_owned()
        })
        .into_owned()
}

/// `EscapeFn` that do not change any thing. Useful when using in a non-html
/// environment.
pub fn no_escape(data: &str) -> String {
    data.to_owned()
}

/// The single entry point of your Handlebars templates
///
/// It maintains compiled templates and registered helpers.
pub struct Registry {
    templates: HashMap<String, Template>,
    helpers: HashMap<String, Box<HelperDef + 'static>>,
    directives: HashMap<String, Box<DirectiveDef + 'static>>,
    escape_fn: EscapeFn,
    source_map: bool,
    strict_mode: bool,
}

impl Debug for Registry {
    fn fmt(&self, f: &mut Formatter) -> Result<(), fmt::Error> {
        f.debug_struct("Handlebars")
            .field("templates", &self.templates)
            .field("helpers", &self.helpers.keys())
            .field("directives", &self.directives.keys())
            .field("source_map", &self.source_map)
            .finish()
    }
}

impl Registry {
    pub fn new() -> Registry {
        let r = Registry {
            templates: HashMap::new(),
            helpers: HashMap::new(),
            directives: HashMap::new(),
            escape_fn: Box::new(html_escape),
            source_map: true,
            strict_mode: false,
        };

        r.setup_builtins()
    }

    fn setup_builtins(mut self) -> Registry {
        self.register_helper("if", Box::new(helpers::IF_HELPER));
        self.register_helper("unless", Box::new(helpers::UNLESS_HELPER));
        self.register_helper("each", Box::new(helpers::EACH_HELPER));
        self.register_helper("with", Box::new(helpers::WITH_HELPER));
        self.register_helper("lookup", Box::new(helpers::LOOKUP_HELPER));
        self.register_helper("raw", Box::new(helpers::RAW_HELPER));
        self.register_helper("log", Box::new(helpers::LOG_HELPER));

        self.register_decorator("inline", Box::new(directives::INLINE_DIRECTIVE));
        self
    }

    /// Enable handlebars template source map
    ///
    /// Source map provides line/col reporting on error. It uses slightly
    /// more memory to maintain the data.
    ///
    /// Default is true.
    pub fn source_map_enabled(&mut self, enable: bool) {
        self.source_map = enable;
    }

    /// Enable handlebars strict mode
    ///
    /// By default, handlebars renders empty string for value that
    /// undefined or never exists. Since rust is a static type
    /// language, we offer strict mode in handlebars-rust.  In strict
    /// mode, if you were access a value that doesn't exist, a
    /// `RenderError` will be raised.
    pub fn set_strict_mode(&mut self, enable: bool) {
        self.strict_mode = enable;
    }

    /// Return strict mode state, default is false.
    ///
    /// By default, handlebars renders empty string for value that
    /// undefined or never exists. Since rust is a static type
    /// language, we offer strict mode in handlebars-rust.  In strict
    /// mode, if you were access a value that doesn't exist, a
    /// `RenderError` will be raised.
    pub fn strict_mode(&self) -> bool {
        self.strict_mode
    }

    /// Register a template string
    ///
    /// Returns `TemplateError` if there is syntax error on parsing template.
    pub fn register_template_string<S>(
        &mut self,
        name: &str,
        tpl_str: S,
    ) -> Result<(), TemplateError>
    where
        S: AsRef<str>,
    {
        try!(
            Template::compile_with_name(tpl_str, name.to_owned(), self.source_map)
                .and_then(|t| Ok(self.templates.insert(name.to_string(), t)))
        );
        Ok(())
    }

    /// Register a partial string
    ///
    /// A named partial will be added to the registry. It will overwrite template with
    /// same name. Currently registered partial is just identical to template.
    pub fn register_partial<S>(&mut self, name: &str, partial_str: S) -> Result<(), TemplateError>
    where
        S: AsRef<str>,
    {
        self.register_template_string(name, partial_str)
    }

    /// Register a template from a path
    pub fn register_template_file
(
        &mut self,
        name: &str,
        tpl_path: P,
    ) -> Result<(), TemplateFileError>
    where
        P: AsRef<Path>,
    {
        let mut file =
            try!(File::open(tpl_path).map_err(|e| TemplateFileError::IOError(e, name.to_owned())));
        self.register_template_source(name, &mut file)
    }

    /// Register a template from `std::io::Read` source
    pub fn register_template_source(
        &mut self,
        name: &str,
        tpl_source: &mut Read,
    ) -> Result<(), TemplateFileError> {
        let mut buf = String::new();
        try!(
            tpl_source
                .read_to_string(&mut buf)
                .map_err(|e| TemplateFileError::IOError(e, name.to_owned()))
        );
        try!(self.register_template_string(name, buf));
        Ok(())
    }

    /// remove a template from the registry
    pub fn unregister_template(&mut self, name: &str) {
        self.templates.remove(name);
    }

    /// register a helper
    pub fn register_helper(
        &mut self,
        name: &str,
        def: Box<HelperDef + 'static>,
    ) -> Option<Box<HelperDef + 'static>> {
        self.helpers.insert(name.to_string(), def)
    }

    /// register a decorator
    pub fn register_decorator(
        &mut self,
        name: &str,
        def: Box<DirectiveDef + 'static>,
    ) -> Option<Box<DirectiveDef + 'static>> {
        self.directives.insert(name.to_string(), def)
    }

    /// Register a new *escape fn* to be used from now on by this registry.
    pub fn register_escape_fn<F: 'static + Fn(&str) -> String + Send + Sync>(
        &mut self,
        escape_fn: F,
    ) {
        self.escape_fn = Box::new(escape_fn);
    }

    /// Restore the default *escape fn*.
    pub fn unregister_escape_fn(&mut self) {
        self.escape_fn = Box::new(html_escape);
    }

    /// Get a reference to the current *escape fn*.
    pub fn get_escape_fn(&self) -> &Fn(&str) -> String {
        &*self.escape_fn
    }

    /// Return a registered template,
    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    /// Return a registered helper
    pub fn get_helper(&self, name: &str) -> Option<&Box<HelperDef + 'static>> {
        self.helpers.get(name)
    }

    /// Return a registered directive, aka decorator
    pub fn get_decorator(&self, name: &str) -> Option<&Box<DirectiveDef + 'static>> {
        self.directives.get(name)
    }

    /// Return all templates registered
    pub fn get_templates(&self) -> &HashMap<String, Template> {
        &self.templates
    }

    /// Unregister all templates
    pub fn clear_templates(&mut self) {
        self.templates.clear();
    }

    /// Render a registered template with some data into a string
    ///
    /// * `name` is the template name you registred previously
    /// * `ctx` is the data that implements `serde::Serialize`
    ///
    /// Returns rendered string or an struct with error information
    pub fn render<T>(&self, name: &str, data: &T) -> Result<String, RenderError>
    where
        T: Serialize,
    {
        let mut writer = StringWriter::new();
        {
            try!(self.render_to_write(name, data, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// Render a registered template and write some data to the `std::io::Write`
    pub fn render_to_write<T>(
        &self,
        name: &str,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), RenderError>
    where
        T: Serialize,
    {
        self.get_template(&name.to_string())
            .ok_or(RenderError::new(format!("Template not found: {}", name)))
            .and_then(|t| {
                let ctx = try!(Context::wraps(data));
                let mut local_helpers = HashMap::new();
                let mut render_context = RenderContext::new(ctx, &mut local_helpers, writer);
                render_context.root_template = t.name.clone();
                t.render(self, &mut render_context)
            })
    }

    /// render a template string using current registry without register it
    pub fn render_template<T>(
        &self,
        template_string: &str,
        data: &T,
    ) -> Result<String, TemplateRenderError>
    where
        T: Serialize,
    {
        let mut writer = StringWriter::new();
        {
            try!(self.render_template_to_write(template_string, data, &mut writer));
        }
        Ok(writer.to_string())
    }

    /// render a template string using current registry without register it
    pub fn render_template_to_write<T>(
        &self,
        template_string: &str,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        let tpl = try!(Template::compile2(template_string, self.source_map));
        let ctx = try!(Context::wraps(data));
        let mut local_helpers = HashMap::new();
        let mut render_context = RenderContext::new(ctx, &mut local_helpers, writer);
        tpl.render(self, &mut render_context)
            .map_err(TemplateRenderError::from)
    }

    /// render a template source using current registry without register it
    pub fn render_template_source_to_write<T>(
        &self,
        template_source: &mut Read,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        let mut tpl_str = String::new();
        try!(
            template_source
                .read_to_string(&mut tpl_str)
                .map_err(|e| TemplateRenderError::IOError(e, "Unamed template source".to_owned()))
        );
        self.render_template_to_write(&tpl_str, data, writer)
    }

    #[deprecated(since = "0.30.0", note = "Please use render_to_write instead.")]
    pub fn renderw<T>(&self, name: &str, data: &T, writer: &mut Write) -> Result<(), RenderError>
    where
        T: Serialize,
    {
        self.render_to_write(name, data, writer)
    }

    #[deprecated(since = "0.30.0", note = "Please use render_template instead.")]
    pub fn template_render<T>(
        &self,
        template_string: &str,
        data: &T,
    ) -> Result<String, TemplateRenderError>
    where
        T: Serialize,
    {
        self.render_template(template_string, data)
    }

    #[deprecated(since = "0.30.0", note = "Please use render_template_to_write instead.")]
    pub fn template_renderw<T>(
        &self,
        template_string: &str,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        self.render_template_to_write(template_string, data, writer)
    }

    #[deprecated(since = "0.30.0", note = "Please use render_template_source_to_write instead.")]
    pub fn template_renderw2<T>(
        &self,
        template_source: &mut Read,
        data: &T,
        writer: &mut Write,
    ) -> Result<(), TemplateRenderError>
    where
        T: Serialize,
    {
        self.render_template_source_to_write(template_source, data, writer)
    }
}

#[cfg(test)]
mod test {
    use registry::Registry;
    use render::{Helper, RenderContext, Renderable};
    use helpers::HelperDef;
    use support::str::StringWriter;
    use error::RenderError;

    #[derive(Clone, Copy)]
    struct DummyHelper;

    impl HelperDef for DummyHelper {
        fn call(
            &self,
            h: &Helper,
            r: &Registry,
            rc: &mut RenderContext,
        ) -> Result<(), RenderError> {
            try!(h.template().unwrap().render(r, rc));
            Ok(())
        }
    }

    static DUMMY_HELPER: DummyHelper = DummyHelper;

    #[test]
    fn test_registry_operations() {
        let mut r = Registry::new();

        assert!(r.register_template_string("index", "

").is_ok());
        assert!(r.register_template_string("index2", "

").is_ok());

        assert_eq!(r.templates.len(), 2);

        r.unregister_template("index");
        assert_eq!(r.templates.len(), 1);

        r.clear_templates();
        assert_eq!(r.templates.len(), 0);

        r.register_helper("dummy", Box::new(DUMMY_HELPER));

        // built-in helpers plus 1
        assert_eq!(r.helpers.len(), 7 + 1);
    }

    #[test]
    fn test_render_to_write() {
        let mut r = Registry::new();

        assert!(r.register_template_string("index", "

").is_ok());

        let mut sw = StringWriter::new();
        {
            r.render_to_write("index", &(), &mut sw).ok().unwrap();
        }

        assert_eq!("

".to_string(), sw.to_string());
    }

    #[test]
    fn test_escape_fn() {
        let mut r = Registry::new();

        let input = String::from("\"<>&");

        r.register_template_string("test", String::from("{{this}}"))
            .unwrap();

        assert_eq!(""<>&", r.render("test", &input).unwrap());

        r.register_escape_fn(|s| s.into());

        assert_eq!("\"<>&", r.render("test", &input).unwrap());

        r.unregister_escape_fn();

        assert_eq!(""<>&", r.render("test", &input).unwrap());
    }

    #[test]
    fn test_escape() {
        let r = Registry::new();
        let data = json!({
            "hello": "world"
        });

        assert_eq!(
            "{{hello}}",
            r.render_template(r"\\{{hello}}", &data).unwrap()
        );
    }

    #[test]
    fn test_strict_mode() {
        use error::TemplateRenderError;
        let mut r = Registry::new();
        assert!(!r.strict_mode());

        r.set_strict_mode(true);
        assert!(r.strict_mode());

        let data = json!({
            "the_only_key": "the_only_value"
        });

        assert!(
            r.render_template("accessing the_only_key {{the_only_key}}", &data)
                .is_ok()
        );
        assert!(
            r.render_template("accessing non-exists key {{the_key_never_exists}}", &data)
                .is_err()
        );

        let render_error =
            r.render_template("accessing non-exists key {{the_key_never_exists}}", &data)
                .unwrap_err();
        if let TemplateRenderError::RenderError(e) = render_error {
            assert_eq!(e.column_no.unwrap(), 26);
        } else {
            unreachable!();
        }
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;
use std::sync::RwLock;
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::path::Path;
use serialize::Encodable;
use hyper::status::StatusCode;
use hyper::server::Response as HyperResponse;
use hyper::header::{
    Headers, Date, HttpDate, Server, ContentType, ContentLength, Header, HeaderFormat
};
use hyper::net::{Fresh, Streaming};
use time;
use mimes::MediaType;
use mustache;
use mustache::Template;
use std::io::{self, Read, Write, copy};
use std::fs::File;
use std::any::Any;
use {NickelError, Halt, MiddlewareResult, Responder};
use modifier::Modifier;

pub type TemplateCache = RwLock<HashMap<String, Template>>;

///A container for the response
pub struct Response<'a, D: 'a = (), T: 'static + Any = Fresh> {
    ///the original `hyper::server::Response`
    origin: HyperResponse<'a, T>,
    templates: &'a TemplateCache,
    data: &'a D,
}

impl<'a, D> Response<'a, D, Fresh> {
    pub fn from_internal<'c, 'd>(response: HyperResponse<'c, Fresh>,
                                 templates: &'c TemplateCache,
                                 data: &'c D)
                                -> Response<'c, D, Fresh> {
        Response {
            origin: response,
            templates: templates,
            data: data
        }
    }

    /// Get a mutable reference to the status.
    pub fn status_mut(&mut self) -> &mut StatusCode {
        self.origin.status_mut()
    }

    /// Get a mutable reference to the Headers.
    pub fn headers_mut(&mut self) -> &mut Headers {
        self.origin.headers_mut()
    }

    /// Modify the response with the provided data.
    ///
    /// # Examples
    /// ```{rust}
    /// extern crate hyper;
    /// #[macro_use] extern crate nickel;
    ///
    /// use nickel::{Nickel, HttpRouter, MediaType};
    /// use nickel::status::StatusCode;
    /// use hyper::header::Location;
    ///
    /// fn main() {
    ///     let mut server = Nickel::new();
    ///     server.get("/a", middleware! { |_, mut res|
    ///             // set the Status
    ///         res.set(StatusCode::PermanentRedirect)
    ///             // update a Header value
    ///            .set(Location("http://nickel.rs".into()));
    ///
    ///         ""
    ///     });
    ///
    ///     server.get("/b", middleware! { |_, mut res|
    ///             // setting the content type
    ///         res.set(MediaType::Json);
    ///
    ///         "{'foo': 'bar'}"
    ///     });
    ///
    ///     // ...
    /// }
    /// ```
    pub fn set<T: Modifier<Response<'a, D>>>(&mut self, attribute: T) -> &mut Response<'a, D> {
        attribute.modify(self);
        self
    }

    /// Writes a response
    ///
    /// # Examples
    /// ```{rust}
    /// use nickel::{Request, Response, MiddlewareResult};
    ///
    /// # #[allow(dead_code)]
    /// fn handler<'a, D>(_: &mut Request<D>, res: Response<'a, D>) -> MiddlewareResult<'a, D> {
    ///     res.send("hello world")
    /// }
    /// ```
    #[inline]
    pub fn send<T: Responder<D>>(self, data: T) -> MiddlewareResult<'a, D> {
        data.respond(self)
    }

    /// Writes a file to the output.
    ///
    /// # Examples
    /// ```{rust}
    /// use nickel::{Request, Response, MiddlewareResult};
    /// use std::path::Path;
    ///
    /// # #[allow(dead_code)]
    /// fn handler<'a, D>(_: &mut Request<D>, res: Response<'a, D>) -> MiddlewareResult<'a, D> {
    ///     let favicon = Path::new("/assets/favicon.ico");
    ///     res.send_file(favicon)
    /// }
    /// ```
    pub fn send_file<P:AsRef<Path>>(mut self, path: P) -> MiddlewareResult<'a, D> {
        let path = path.as_ref();
        // Chunk the response
        self.origin.headers_mut().remove::<ContentLength>();
        // Determine content type by file extension or default to binary
        let mime = mime_from_filename(path).unwrap_or(MediaType::Bin);
        self.set(mime);

        let mut file = try_with!(self, {
            File::open(path).map_err(|e| format!("Failed to send file '{:?}': {}",
                                                 path, e))
        });

        let mut stream = try!(self.start());
        match copy(&mut file, &mut stream) {
            Ok(_) => Ok(Halt(stream)),
            Err(e) => stream.bail(format!("Failed to send file: {}", e))
        }
    }

    // TODO: This needs to be more sophisticated to return the correct headers
    // not just "some headers" :)
    //
    // Also, it should only set them if not already set.
    fn set_fallback_headers(&mut self) {
        self.set_header_fallback(|| Date(HttpDate(time::now_utc())));
        self.set_header_fallback(|| Server("Nickel".to_string()));
        self.set_header_fallback(|| ContentType(MediaType::Html.into()));
    }

    /// Return an error with the appropriate status code for error handlers to
    /// provide output for.
    pub fn error<T>(self, status: StatusCode, message: T) -> MiddlewareResult<'a, D>
            where T: Into<Cow<'static, str>> {
        Err(NickelError::new(self, message, status))
    }

    /// Sets the header if not already set.
    ///
    /// If the header is not set then `f` will be called.
    /// Renders the given template bound with the given data.
    ///
    /// # Examples
    /// ```{rust}
    /// #[macro_use] extern crate nickel;
    /// extern crate hyper;
    ///
    /// use nickel::{Nickel, HttpRouter, MediaType};
    /// use hyper::header::ContentType;
    ///
    /// # #[allow(unreachable_code)]
    /// fn main() {
    ///     let mut server = Nickel::new();
    ///     server.get("/", middleware! { |_, mut res|
    ///         res.set(MediaType::Html);
    ///         res.set_header_fallback(|| {
    ///             panic!("Should not get called");
    ///             ContentType(MediaType::Txt.into())
    ///         });
    ///
    ///         "<h1>Hello World</h1>"
    ///     });
    ///
    ///     // ...
    /// }
    /// ```
    pub fn set_header_fallback<F, H>(&mut self, f: F)
            where H: Header + HeaderFormat, F: FnOnce() -> H {
        let headers = self.origin.headers_mut();
        if !headers.has::<H>() { headers.set(f()) }
    }

    /// Renders the given template bound with the given data.
    ///
    /// # Examples
    /// ```{rust}
    /// use std::collections::HashMap;
    /// use nickel::{Request, Response, MiddlewareResult};
    ///
    /// # #[allow(dead_code)]
    /// fn handler<'a, D>(_: &mut Request<D>, res: Response<'a, D>) -> MiddlewareResult<'a, D> {
    ///     let mut data = HashMap::new();
    ///     data.insert("name", "user");
    ///     res.render("examples/assets/template.tpl", &data)
    /// }
    /// ```
    pub fn render<T, P>(self, path: P, data: &T) -> MiddlewareResult<'a, D>
            where T: Encodable, P: AsRef<str> + Into<String> {
        fn render<'a, D, T>(res: Response<'a, D>, template: &Template, data: &T)
                -> MiddlewareResult<'a, D> where T: Encodable {
            let mut stream = try!(res.start());
            match template.render(&mut stream, data) {
                Ok(()) => Ok(Halt(stream)),
                Err(e) => stream.bail(format!("Problem rendering template: {:?}", e))
            }
        }

        // Fast path doesn't need writer lock
        if let Some(t) = self.templates.read().unwrap().get(path.as_ref()) {
            return render(self, t, data);
        }

        // We didn't find the template, get writers lock
        let mut templates = self.templates.write().unwrap();

        // Additional clone required for now as the entry api doesn't give us a key ref
        let path = path.into();

        // Search again incase there was a race to compile the template
        let template = match templates.entry(path.clone()) {
            Vacant(entry) => {
                let template = try_with!(self, {
                    mustache::compile_path(&path)
                             .map_err(|e| format!("Failed to compile template '{}': {:?}",
                                            path, e))
                });
                entry.insert(template)
            },
            Occupied(entry) => entry.into_mut()
        };

        render(self, template, data)
    }

    pub fn start(mut self) -> Result<Response<'a, D, Streaming>, NickelError<'a, D>> {
        self.set_fallback_headers();

        let Response { origin, templates, data } = self;
        match origin.start() {
            Ok(origin) => {
                Ok(Response {
                    origin: origin,
                    templates: templates,
                    data: data
                })
            },
            Err(e) =>
                unsafe {
                    Err(NickelError::without_response(format!("Failed to start response: {}", e)))
                }
        }
    }

    pub fn server_data(&self) -> &D {
        &self.data
    }
}

impl<'a, 'b, D> Write for Response<'a, D, Streaming> {
    #[inline(always)]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.origin.write(buf)
    }

    #[inline(always)]
    fn flush(&mut self) -> io::Result<()> {
        self.origin.flush()
    }
}

impl<'a, 'b, D> Response<'a, D, Streaming> {
    /// In the case of an unrecoverable error while a stream is already in
    /// progress, there is no standard way to signal to the client that an
    /// error has occurred. `bail` will drop the connection and log an error
    /// message.
    pub fn bail<T>(self, message: T) -> MiddlewareResult<'a, D>
            where T: Into<Cow<'static, str>> {
        let _ = self.end();
        unsafe { Err(NickelError::without_response(message)) }
    }

    /// Flushes all writing of a response to the client.
    pub fn end(self) -> io::Result<()> {
        self.origin.end()
    }
}

impl <'a, D, T: 'static + Any> Response<'a, D, T> {
    /// The status of this response.
    pub fn status(&self) -> StatusCode {
        self.origin.status()
    }

    /// The headers of this response.
    pub fn headers(&self) -> &Headers {
        self.origin.headers()
    }
}

fn mime_from_filename<P: AsRef<Path>>(path: P) -> Option<MediaType> {
    path.as_ref()
        .extension()
        .and_then(|os| os.to_str())
        // Lookup mime from file extension
        .and_then(|s| s.parse().ok())
}

#[test]
fn matches_content_type () {
    assert_eq!(Some(MediaType::Txt), mime_from_filename("test.txt"));
    assert_eq!(Some(MediaType::Json), mime_from_filename("test.json"));
    assert_eq!(Some(MediaType::Bin), mime_from_filename("test.bin"));
}

mod modifier_impls {
    use hyper::header::*;
    use hyper::status::StatusCode;
    use modifier::Modifier;
    use {Response, MediaType};

    impl<'a, D> Modifier<Response<'a, D>> for StatusCode {
        fn modify(self, res: &mut Response<'a, D>) {
            *res.status_mut() = self
        }
    }

    impl<'a, D> Modifier<Response<'a, D>> for MediaType {
        fn modify(self, res: &mut Response<'a, D>) {
            ContentType(self.into()).modify(res)
        }
    }

    macro_rules! header_modifiers {
        ($($t:ty),+) => (
            $(
                impl<'a, D> Modifier<Response<'a, D>> for $t {
                    fn modify(self, res: &mut Response<'a, D>) {
                        res.headers_mut().set(self)
                    }
                }
            )+
        )
    }

    header_modifiers! {
        Accept,
        AccessControlAllowHeaders,
        AccessControlAllowMethods,
        AccessControlAllowOrigin,
        AccessControlMaxAge,
        AccessControlRequestHeaders,
        AccessControlRequestMethod,
        AcceptCharset,
        AcceptEncoding,
        AcceptLanguage,
        AcceptRanges,
        Allow,
        Authorization<Basic>,
        Authorization<String>,
        CacheControl,
        Cookie,
        Connection,
        ContentEncoding,
        ContentLanguage,
        ContentLength,
        ContentType,
        Date,
        ETag,
        Expect,
        Expires,
        From,
        Host,
        IfMatch,
        IfModifiedSince,
        IfNoneMatch,
        IfRange,
        IfUnmodifiedSince,
        LastModified,
        Location,
        Pragma,
        Referer,
        Server,
        SetCookie,
        TransferEncoding,
        Upgrade,
        UserAgent,
        Vary
    }
}
<commit_msg>feat(response): allow Plugins for Response<commit_after>use std::mem;
use std::borrow::Cow;
use std::sync::RwLock;
use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied, Vacant};
use std::path::Path;
use serialize::Encodable;
use hyper::status::StatusCode;
use hyper::server::Response as HyperResponse;
use hyper::header::{
    Headers, Date, HttpDate, Server, ContentType, ContentLength, Header, HeaderFormat
};
use hyper::net::{Fresh, Streaming};
use time;
use mimes::MediaType;
use mustache;
use mustache::Template;
use std::io::{self, Read, Write, copy};
use std::fs::File;
use std::any::Any;
use {NickelError, Halt, MiddlewareResult, Responder};
use modifier::Modifier;
use plugin::{Extensible, Pluggable};
use typemap::TypeMap;

pub type TemplateCache = RwLock<HashMap<String, Template>>;

///A container for the response
pub struct Response<'a, D: 'a = (), T: 'static + Any = Fresh> {
    ///the original `hyper::server::Response`
    origin: HyperResponse<'a, T>,
    templates: &'a TemplateCache,
    data: &'a D,
    map: TypeMap,
    // This should be FnBox, but that's currently unstable
    on_send: Vec<Box<FnMut(&mut Response<'a, D, Fresh>)>>
}

impl<'a, D> Response<'a, D, Fresh> {
    pub fn from_internal<'c, 'd>(response: HyperResponse<'c, Fresh>,
                                 templates: &'c TemplateCache,
                                 data: &'c D)
                                -> Response<'c, D, Fresh> {
        Response {
            origin: response,
            templates: templates,
            data: data,
            map: TypeMap::new(),
            on_send: vec![]
        }
    }

    /// Get a mutable reference to the status.
    pub fn status_mut(&mut self) -> &mut StatusCode {
        self.origin.status_mut()
    }

    /// Get a mutable reference to the Headers.
    pub fn headers_mut(&mut self) -> &mut Headers {
        self.origin.headers_mut()
    }

    /// Modify the response with the provided data.
    ///
    /// # Examples
    /// ```{rust}
    /// extern crate hyper;
    /// #[macro_use] extern crate nickel;
    ///
    /// use nickel::{Nickel, HttpRouter, MediaType};
    /// use nickel::status::StatusCode;
    /// use hyper::header::Location;
    ///
    /// fn main() {
    ///     let mut server = Nickel::new();
    ///     server.get("/a", middleware! { |_, mut res|
    ///             // set the Status
    ///         res.set(StatusCode::PermanentRedirect)
    ///             // update a Header value
    ///            .set(Location("http://nickel.rs".into()));
    ///
    ///         ""
    ///     });
    ///
    ///     server.get("/b", middleware! { |_, mut res|
    ///             // setting the content type
    ///         res.set(MediaType::Json);
    ///
    ///         "{'foo': 'bar'}"
    ///     });
    ///
    ///     // ...
    /// }
    /// ```
    pub fn set<T: Modifier<Response<'a, D>>>(&mut self, attribute: T) -> &mut Response<'a, D> {
        attribute.modify(self);
        self
    }

    /// Writes a response
    ///
    /// # Examples
    /// ```{rust}
    /// use nickel::{Request, Response, MiddlewareResult};
    ///
    /// # #[allow(dead_code)]
    /// fn handler<'a, D>(_: &mut Request<D>, res: Response<'a, D>) -> MiddlewareResult<'a, D> {
    ///     res.send("hello world")
    /// }
    /// ```
    #[inline]
    pub fn send<T: Responder<D>>(self, data: T) -> MiddlewareResult<'a, D> {
        data.respond(self)
    }

    /// Writes a file to the output.
    ///
    /// # Examples
    /// ```{rust}
    /// use nickel::{Request, Response, MiddlewareResult};
    /// use std::path::Path;
    ///
    /// # #[allow(dead_code)]
    /// fn handler<'a, D>(_: &mut Request<D>, res: Response<'a, D>) -> MiddlewareResult<'a, D> {
    ///     let favicon = Path::new("/assets/favicon.ico");
    ///     res.send_file(favicon)
    /// }
    /// ```
    pub fn send_file<P:AsRef<Path>>(mut self, path: P) -> MiddlewareResult<'a, D> {
        let path = path.as_ref();
        // Chunk the response
        self.origin.headers_mut().remove::<ContentLength>();
        // Determine content type by file extension or default to binary
        let mime = mime_from_filename(path).unwrap_or(MediaType::Bin);
        self.set(mime);

        let mut file = try_with!(self, {
            File::open(path).map_err(|e| format!("Failed to send file '{:?}': {}",
                                                 path, e))
        });

        let mut stream = try!(self.start());
        match copy(&mut file, &mut stream) {
            Ok(_) => Ok(Halt(stream)),
            Err(e) => stream.bail(format!("Failed to send file: {}", e))
        }
    }

    // TODO: This needs to be more sophisticated to return the correct headers
    // not just "some headers" :)
    //
    // Also, it should only set them if not already set.
    fn set_fallback_headers(&mut self) {
        self.set_header_fallback(|| Date(HttpDate(time::now_utc())));
        self.set_header_fallback(|| Server("Nickel".to_string()));
        self.set_header_fallback(|| ContentType(MediaType::Html.into()));
    }

    /// Return an error with the appropriate status code for error handlers to
    /// provide output for.
    pub fn error<T>(self, status: StatusCode, message: T) -> MiddlewareResult<'a, D>
            where T: Into<Cow<'static, str>> {
        Err(NickelError::new(self, message, status))
    }

    /// Sets the header if not already set.
    ///
    /// If the header is not set then `f` will be called.
    /// Renders the given template bound with the given data.
    ///
    /// # Examples
    /// ```{rust}
    /// #[macro_use] extern crate nickel;
    /// extern crate hyper;
    ///
    /// use nickel::{Nickel, HttpRouter, MediaType};
    /// use hyper::header::ContentType;
    ///
    /// # #[allow(unreachable_code)]
    /// fn main() {
    ///     let mut server = Nickel::new();
    ///     server.get("/", middleware! { |_, mut res|
    ///         res.set(MediaType::Html);
    ///         res.set_header_fallback(|| {
    ///             panic!("Should not get called");
    ///             ContentType(MediaType::Txt.into())
    ///         });
    ///
    ///         "<h1>Hello World</h1>"
    ///     });
    ///
    ///     // ...
    /// }
    /// ```
    pub fn set_header_fallback<F, H>(&mut self, f: F)
            where H: Header + HeaderFormat, F: FnOnce() -> H {
        let headers = self.origin.headers_mut();
        if !headers.has::<H>() { headers.set(f()) }
    }

    /// Renders the given template bound with the given data.
    ///
    /// # Examples
    /// ```{rust}
    /// use std::collections::HashMap;
    /// use nickel::{Request, Response, MiddlewareResult};
    ///
    /// # #[allow(dead_code)]
    /// fn handler<'a, D>(_: &mut Request<D>, res: Response<'a, D>) -> MiddlewareResult<'a, D> {
    ///     let mut data = HashMap::new();
    ///     data.insert("name", "user");
    ///     res.render("examples/assets/template.tpl", &data)
    /// }
    /// ```
    pub fn render<T, P>(self, path: P, data: &T) -> MiddlewareResult<'a, D>
            where T: Encodable, P: AsRef<str> + Into<String> {
        fn render<'a, D, T>(res: Response<'a, D>, template: &Template, data: &T)
                -> MiddlewareResult<'a, D> where T: Encodable {
            let mut stream = try!(res.start());
            match template.render(&mut stream, data) {
                Ok(()) => Ok(Halt(stream)),
                Err(e) => stream.bail(format!("Problem rendering template: {:?}", e))
            }
        }

        // Fast path doesn't need writer lock
        if let Some(t) = self.templates.read().unwrap().get(path.as_ref()) {
            return render(self, t, data);
        }

        // We didn't find the template, get writers lock
        let mut templates = self.templates.write().unwrap();

        // Additional clone required for now as the entry api doesn't give us a key ref
        let path = path.into();

        // Search again incase there was a race to compile the template
        let template = match templates.entry(path.clone()) {
            Vacant(entry) => {
                let template = try_with!(self, {
                    mustache::compile_path(&path)
                             .map_err(|e| format!("Failed to compile template '{}': {:?}",
                                            path, e))
                });
                entry.insert(template)
            },
            Occupied(entry) => entry.into_mut()
        };

        render(self, template, data)
    }

    pub fn start(mut self) -> Result<Response<'a, D, Streaming>, NickelError<'a, D>> {
        let on_send = mem::replace(&mut self.on_send, vec![]);
        for mut f in on_send.into_iter().rev() {
            // TODO: Ensure `f` doesn't call on_send again
            f(&mut self)
        }

        // Set fallback headers last after everything runs, if we did this before as an
        // on_send then it would possibly set redundant things.
        self.set_fallback_headers();

        let Response { origin, templates, data, map, on_send } = self;
        match origin.start() {
            Ok(origin) => {
                Ok(Response {
                    origin: origin,
                    templates: templates,
                    data: data,
                    map: map,
                    on_send: on_send
                })
            },
            Err(e) =>
                unsafe {
                    Err(NickelError::without_response(format!("Failed to start response: {}", e)))
                }
        }
    }

    pub fn server_data(&self) -> &D {
        &self.data
    }

    pub fn on_send<F>(&mut self, f: F)
            where F: FnMut(&mut Response<'a, D, Fresh>) + 'static {
        self.on_send.push(Box::new(f))
    }
}

impl<'a, 'b, D> Write for Response<'a, D, Streaming> {
    #[inline(always)]
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.origin.write(buf)
    }

    #[inline(always)]
    fn flush(&mut self) -> io::Result<()> {
        self.origin.flush()
    }
}

impl<'a, 'b, D> Response<'a, D, Streaming> {
    /// In the case of an unrecoverable error while a stream is already in
    /// progress, there is no standard way to signal to the client that an
    /// error has occurred. `bail` will drop the connection and log an error
    /// message.
    pub fn bail<T>(self, message: T) -> MiddlewareResult<'a, D>
            where T: Into<Cow<'static, str>> {
        let _ = self.end();
        unsafe { Err(NickelError::without_response(message)) }
    }

    /// Flushes all writing of a response to the client.
    pub fn end(self) -> io::Result<()> {
        self.origin.end()
    }
}

impl <'a, D, T: 'static + Any> Response<'a, D, T> {
    /// The status of this response.
    pub fn status(&self) -> StatusCode {
        self.origin.status()
    }

    /// The headers of this response.
    pub fn headers(&self) -> &Headers {
        self.origin.headers()
    }

    pub fn data(&self) -> &D {
        &self.data
    }
}

impl<'a, D, T: 'static + Any> Extensible for Response<'a, D, T> {
    fn extensions(&self) -> &TypeMap {
        &self.map
    }

    fn extensions_mut(&mut self) -> &mut TypeMap {
        &mut self.map
    }
}

impl<'a, D, T: 'static + Any> Pluggable for Response<'a, D, T> {}

fn mime_from_filename<P: AsRef<Path>>(path: P) -> Option<MediaType> {
    path.as_ref()
        .extension()
        .and_then(|os| os.to_str())
        // Lookup mime from file extension
        .and_then(|s| s.parse().ok())
}

#[test]
fn matches_content_type () {
    assert_eq!(Some(MediaType::Txt), mime_from_filename("test.txt"));
    assert_eq!(Some(MediaType::Json), mime_from_filename("test.json"));
    assert_eq!(Some(MediaType::Bin), mime_from_filename("test.bin"));
}

mod modifier_impls {
    use hyper::header::*;
    use hyper::status::StatusCode;
    use modifier::Modifier;
    use {Response, MediaType};

    impl<'a, D> Modifier<Response<'a, D>> for StatusCode {
        fn modify(self, res: &mut Response<'a, D>) {
            *res.status_mut() = self
        }
    }

    impl<'a, D> Modifier<Response<'a, D>> for MediaType {
        fn modify(self, res: &mut Response<'a, D>) {
            ContentType(self.into()).modify(res)
        }
    }

    macro_rules! header_modifiers {
        ($($t:ty),+) => (
            $(
                impl<'a, D> Modifier<Response<'a, D>> for $t {
                    fn modify(self, res: &mut Response<'a, D>) {
                        res.headers_mut().set(self)
                    }
                }
            )+
        )
    }

    header_modifiers! {
        Accept,
        AccessControlAllowHeaders,
        AccessControlAllowMethods,
        AccessControlAllowOrigin,
        AccessControlMaxAge,
        AccessControlRequestHeaders,
        AccessControlRequestMethod,
        AcceptCharset,
        AcceptEncoding,
        AcceptLanguage,
        AcceptRanges,
        Allow,
        Authorization<Basic>,
        Authorization<String>,
        CacheControl,
        Cookie,
        Connection,
        ContentEncoding,
        ContentLanguage,
        ContentLength,
        ContentType,
        Date,
        ETag,
        Expect,
        Expires,
        From,
        Host,
        IfMatch,
        IfModifiedSince,
        IfNoneMatch,
        IfRange,
        IfUnmodifiedSince,
        LastModified,
        Location,
        Pragma,
        Referer,
        Server,
        SetCookie,
        TransferEncoding,
        Upgrade,
        UserAgent,
        Vary
    }
}
<|endoftext|>


<commit_before>#![warn(unused_extern_crates)]

#[macro_use]
extern crate clap;
extern crate env_logger;
extern crate rhq_core as rhq;
extern crate shlex;

use std::borrow::Cow;
use std::env;
use std::path::{Path, PathBuf};
use clap::{App, AppSettings, Arg, ArgMatches, SubCommand};

use rhq::{Query, Result, Workspace};
use rhq::util;
use rhq::vcs::Vcs;

const POSSIBLE_VCS: &[&str] = &["git", "hg", "darcs", "pijul"];

fn main() {
    env_logger::init().expect("failed to initialize env_logger.");
    if let Err(message) = run() {
        println!("failed with: {}", message);
        std::process::exit(1);
    }
}

macro_rules! def_app {
    ($( $name:expr => $t:ident, )*) => {
        fn app<'a, 'b: 'a>() -> App<'a, 'b> {
            app_from_crate!()
                .setting(AppSettings::VersionlessSubcommands)
                .setting(AppSettings::SubcommandRequiredElseHelp)
                $( .subcommand($t::app(SubCommand::with_name($name))) )*
        }

        pub fn run() -> Result<()> {
            let matches = app().get_matches();
            match matches.subcommand() {
                $( ($name, Some(m)) => $t::from_matches(m).run(), )*
                _ => unreachable!(),
            }
        }
    }
}

def_app! {
    "add"        => AddCommand,
    "clone"      => CloneCommand,
    "completion" => CompletionCommand,
    "foreach"    => ForeachCommand,
    "import"     => ImportCommand,
    "list"       => ListCommand,
    "new"        => NewCommand,
    "refresh"    => RefreshCommand,
}

/// subcommand `add`
struct AddCommand {
    paths: Vec<PathBuf>,
    verbose: bool,
}

impl AddCommand {
    fn app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Add existed repositories into management")
            .arg_from_usage("[paths]...      'Location of local repositories'")
            .arg_from_usage("-v, --verbose   'Use verbose output'")
    }

    fn from_matches(m: &ArgMatches) -> AddCommand {
        AddCommand {
            paths: m.values_of("paths")
                .map(|s| s.map(PathBuf::from).collect())
                .unwrap_or_else(|| vec![env::current_dir().expect("env::current_dir()")]),
            verbose: m.is_present("verbose"),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?.verbose_output(self.verbose);
        for path in self.paths {
            workspace.add_repository_if_exists(&path)?;
        }
        workspace.save_cache()?;

        Ok(())
    }
}

/// subcommand `import`
struct ImportCommand {
    roots: Option<Vec<PathBuf>>,
    depth: Option<usize>,
    verbose: bool,
}

impl ImportCommand {
    fn app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Import existed repositories into management")
            .arg_from_usage("[roots]...      'Root directories contains for scanning'")
            .arg_from_usage("--depth=[depth] 'Maximal depth of entries for each base directory'")
            .arg_from_usage("-v, --verbose   'Use verbose output'")
    }

    fn from_matches(m: &ArgMatches) -> ImportCommand {
        ImportCommand {
            roots: m.values_of("roots").map(|s| s.map(PathBuf::from).collect()),
            depth: m.value_of("depth").and_then(|s| s.parse().ok()),
            verbose: m.is_present("verbose"),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?.verbose_output(self.verbose);

        let roots = self.roots
            .unwrap_or_else(|| workspace.config().include_dirs());
        for root in roots {
            workspace.import_repositories(root, self.depth)?;
        }
        workspace.save_cache()?;

        Ok(())
    }
}

/// Subommand `refresh`
struct RefreshCommand {
    verbose: bool,
    sort: bool,
}

impl RefreshCommand {
    fn app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Scan repository list and drop if it is not existed or matches exclude pattern.")
            .arg_from_usage("-v, --verbose 'Use verbose output'")
            .arg_from_usage("-s, --sort    'Sort by path string'")
    }

    fn from_matches(m: &ArgMatches) -> RefreshCommand {
        RefreshCommand {
            verbose: m.is_present("verbose"),
            sort: m.is_present("sort"),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?.verbose_output(self.verbose);
        workspace.drop_invalid_repositories();
        if self.sort {
            workspace.sort_repositories();
        }
        workspace.save_cache()?;
        Ok(())
    }
}

/// Subcommand `new`
struct NewCommand<'a> {
    path: &'a str,
    vcs: Vcs,
    hook: Option<Vec<String>>,
}

impl<'a> NewCommand<'a> {
    fn app<'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Create a new repository and add it into management")
            .arg_from_usage("<path>           'Path of target repository, or URL-like pattern'")
            .arg(Arg::from_usage("--vcs=[vcs] 'Used Version Control System'").possible_values(POSSIBLE_VCS))
            .arg_from_usage("--hook=[hook]    'Post hook after initialization'")
    }

    fn from_matches<'b: 'a>(m: &'b ArgMatches<'a>) -> NewCommand<'a> {
        NewCommand {
            path: m.value_of("path").unwrap(),
            vcs: m.value_of("vcs")
                .and_then(|s| s.parse().ok())
                .unwrap_or(Vcs::Git),
            hook: m.value_of("hook").and_then(|s| shlex::split(s)),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?;

        let path: Cow<Path> = match self.path.parse::<Query>() {
            Ok(query) => workspace.resolve_query(&query)?.into(),
            Err(_) => Path::new(self.path).into(),
        };

        // init
        let repo = workspace.create_empty_repository(&path, self.vcs)?;

        // hook
        match (repo, self.hook) {
            (Some(ref repo), Some(ref hook)) if hook.len() >= 1 => {
                workspace.print(format_args!("[info] Running post hook command...\n"));
                repo.run_command(&hook[0], &hook[1..])?;
            }
            _ => {}
        }

        workspace.save_cache()?;
        Ok(())
    }
}

/// Subcommand `clone`
pub struct CloneCommand<'a> {
    query: Query,
    dest: Option<PathBuf>,
    root: Option<&'a Path>,
    ssh: bool,
    args: Vec<&'a str>,
    vcs: Vcs,
}

impl<'a> CloneCommand<'a> {
    fn app<'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Clone remote repositories, and then add it under management")
            .arg_from_usage("<query>          'an URL or a string to determine the URL of remote repository'")
            .arg_from_usage("[dest]           'Destination directory of cloned repository'")
            .arg_from_usage("[args]...        'Supplemental arguments for VCS command'")
            .arg_from_usage("--root=[root]    'Path to determine the destination directory of cloned repository'")
            .arg_from_usage("-s, --ssh        'Use SSH protocol'")
            .arg(Arg::from_usage("--vcs=[vcs] 'Used Version Control System'").possible_values(POSSIBLE_VCS))
    }

    fn from_matches<'b: 'a>(m: &'b ArgMatches<'a>) -> CloneCommand<'a> {
        CloneCommand {
            query: m.value_of("query").and_then(|s| s.parse().ok()).unwrap(),
            dest: m.value_of("dest").map(PathBuf::from),
            root: m.value_of("root").map(Path::new),
            ssh: m.is_present("ssh"),
            args: m.values_of("args").map(|s| s.collect()).unwrap_or_default(),
            vcs: m.value_of("vcs")
                .and_then(|s| s.parse().ok())
                .unwrap_or(Vcs::Git),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?.root_dir(self.root);

        let url = self.query.to_url(self.ssh)?;
        let dest = match self.dest {
            Some(dest) => dest,
            None => workspace.resolve_query(&self.query)?,
        };
        workspace.clone_repository(&url, &dest, self.vcs, &self.args[..])?;

        workspace.save_cache()?;
        Ok(())
    }
}

#[derive(Debug)]
enum ListFormat {
    Name,
    FullPath,
}

impl ::std::str::FromStr for ListFormat {
    type Err = ();
    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> {
        match s {
            "name" => Ok(ListFormat::Name),
            "fullpath" => Ok(ListFormat::FullPath),
            _ => Err(()),
        }
    }
}

/// Subcommand `list`
pub struct ListCommand {
    format: ListFormat,
}

impl ListCommand {
    fn app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("List local repositories managed by rhq")
            .arg(clap::Arg::from_usage("--format=[format] 'List format'").possible_values(&["name", "fullpath"]))
    }

    fn from_matches(m: &ArgMatches) -> ListCommand {
        ListCommand {
            format: m.value_of("format")
                .and_then(|s| s.parse().ok())
                .unwrap_or(ListFormat::FullPath),
        }
    }

    fn run(self) -> Result<()> {
        let workspace = Workspace::new()?;
        workspace.for_each_repo(|repo| {
            match self.format {
                ListFormat::Name => println!("{}", repo.name()),
                ListFormat::FullPath => println!("{}", repo.path_string()),
            }
            Ok(())
        })
    }
}

/// Subcommand `foreach`
pub struct ForeachCommand<'a> {
    command: &'a str,
    args: Vec<&'a str>,
    dry_run: bool,
}

impl<'a> ForeachCommand<'a> {
    fn app<'b: 'a>(app: App<'a, 'b>) -> clap::App<'a, 'b> {
        app.about("Execute command into each repositories")
            .arg_from_usage("<command>       'Command name'")
            .arg_from_usage("[args]...       'Arguments of command'")
            .arg_from_usage("-n, --dry-run   'Do not actually execute command'")
    }

    fn from_matches<'b: 'a>(m: &'b ArgMatches<'a>) -> ForeachCommand<'a> {
        ForeachCommand {
            command: m.value_of("command").unwrap(),
            args: m.values_of("args").map(|s| s.collect()).unwrap_or_default(),
            dry_run: m.is_present("dry-run"),
        }
    }

    fn run(self) -> Result<()> {
        let workspace = Workspace::new()?;
        workspace.for_each_repo(|repo| {
            if self.dry_run {
                workspace.print(format_args!(
                    "+ {} {}",
                    self.command,
                    util::join_str(&self.args[..])
                ));
            } else {
                repo.run_command(self.command, &self.args[..])?;
            }
            Ok(())
        })
    }
}

pub struct CompletionCommand<'a> {
    shell: clap::Shell,
    out_file: Option<&'a Path>,
}

impl<'a> CompletionCommand<'a> {
    fn app<'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Generate completion scripts for your shell")
            .setting(AppSettings::ArgRequiredElseHelp)
            .arg(
                clap::Arg::with_name("shell")
                    .help("target shell")
                    .possible_values(&["bash", "zsh", "fish", "powershell"])
                    .required(true),
            )
            .arg(clap::Arg::from_usage("[out-file]  'path to output script'"))
    }

    fn from_matches<'b: 'a>(m: &'b ArgMatches<'a>) -> CompletionCommand<'a> {
        CompletionCommand {
            shell: m.value_of("shell").and_then(|s| s.parse().ok()).unwrap(),
            out_file: m.value_of("out-file").map(Path::new),
        }
    }

    fn run(self) -> Result<()> {
        if let Some(path) = self.out_file {
            let mut file = ::std::fs::OpenOptions::new()
                .write(true)
                .create(true)
                .append(false)
                .open(path)
                .unwrap();
            app().gen_completions_to(env!("CARGO_PKG_NAME"), self.shell, &mut file);
        } else {
            app().gen_completions_to(env!("CARGO_PKG_NAME"), self.shell, &mut ::std::io::stdout());
        }
        Ok(())
    }
}
<commit_msg>refactor(cli): add<commit_after>#![warn(unused_extern_crates)]

#[macro_use]
extern crate clap;
extern crate env_logger;
extern crate rhq_core as rhq;
extern crate shlex;

use std::borrow::Cow;
use std::env;
use std::path::{Path, PathBuf};
use clap::{App, AppSettings, Arg, ArgMatches, SubCommand};

use rhq::{Query, Result, Workspace};
use rhq::util;
use rhq::vcs::Vcs;

const POSSIBLE_VCS: &[&str] = &["git", "hg", "darcs", "pijul"];

fn main() {
    env_logger::init().expect("failed to initialize env_logger.");
    if let Err(message) = run() {
        println!("failed with: {}", message);
        std::process::exit(1);
    }
}

macro_rules! def_app {
    ($( $name:expr => $t:ident, )*) => {
        fn app<'a, 'b: 'a>() -> App<'a, 'b> {
            app_from_crate!()
                .setting(AppSettings::VersionlessSubcommands)
                .setting(AppSettings::SubcommandRequiredElseHelp)
                $( .subcommand($t::app(SubCommand::with_name($name))) )*
        }

        pub fn run() -> Result<()> {
            let matches = app().get_matches();
            match matches.subcommand() {
                $( ($name, Some(m)) => $t::from_matches(m).run(), )*
                _ => unreachable!(),
            }
        }
    }
}

def_app! {
    "add"        => AddCommand,
    "clone"      => CloneCommand,
    "completion" => CompletionCommand,
    "foreach"    => ForeachCommand,
    "import"     => ImportCommand,
    "list"       => ListCommand,
    "new"        => NewCommand,
    "refresh"    => RefreshCommand,
}

/// subcommand `add`
struct AddCommand {
    paths: Option<Vec<PathBuf>>,
    verbose: bool,
}

impl AddCommand {
    fn app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Add existed repositories into management")
            .arg_from_usage("[paths]...      'Location of local repositories'")
            .arg_from_usage("-v, --verbose   'Use verbose output'")
    }

    fn from_matches(m: &ArgMatches) -> AddCommand {
        AddCommand {
            paths: m.values_of("paths").map(|s| s.map(PathBuf::from).collect()),
            verbose: m.is_present("verbose"),
        }
    }

    fn run(self) -> Result<()> {
        let paths = self.paths
            .unwrap_or_else(|| vec![env::current_dir().expect("env::current_dir()")]);

        let mut workspace = Workspace::new()?.verbose_output(self.verbose);
        for path in paths {
            workspace.add_repository_if_exists(&path)?;
        }
        workspace.save_cache()?;

        Ok(())
    }
}

/// subcommand `import`
struct ImportCommand {
    roots: Option<Vec<PathBuf>>,
    depth: Option<usize>,
    verbose: bool,
}

impl ImportCommand {
    fn app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Import existed repositories into management")
            .arg_from_usage("[roots]...      'Root directories contains for scanning'")
            .arg_from_usage("--depth=[depth] 'Maximal depth of entries for each base directory'")
            .arg_from_usage("-v, --verbose   'Use verbose output'")
    }

    fn from_matches(m: &ArgMatches) -> ImportCommand {
        ImportCommand {
            roots: m.values_of("roots").map(|s| s.map(PathBuf::from).collect()),
            depth: m.value_of("depth").and_then(|s| s.parse().ok()),
            verbose: m.is_present("verbose"),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?.verbose_output(self.verbose);

        let roots = self.roots
            .unwrap_or_else(|| workspace.config().include_dirs());
        for root in roots {
            workspace.import_repositories(root, self.depth)?;
        }
        workspace.save_cache()?;

        Ok(())
    }
}

/// Subommand `refresh`
struct RefreshCommand {
    verbose: bool,
    sort: bool,
}

impl RefreshCommand {
    fn app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Scan repository list and drop if it is not existed or matches exclude pattern.")
            .arg_from_usage("-v, --verbose 'Use verbose output'")
            .arg_from_usage("-s, --sort    'Sort by path string'")
    }

    fn from_matches(m: &ArgMatches) -> RefreshCommand {
        RefreshCommand {
            verbose: m.is_present("verbose"),
            sort: m.is_present("sort"),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?.verbose_output(self.verbose);
        workspace.drop_invalid_repositories();
        if self.sort {
            workspace.sort_repositories();
        }
        workspace.save_cache()?;
        Ok(())
    }
}

/// Subcommand `new`
struct NewCommand<'a> {
    path: &'a str,
    vcs: Vcs,
    hook: Option<Vec<String>>,
}

impl<'a> NewCommand<'a> {
    fn app<'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Create a new repository and add it into management")
            .arg_from_usage("<path>           'Path of target repository, or URL-like pattern'")
            .arg(Arg::from_usage("--vcs=[vcs] 'Used Version Control System'").possible_values(POSSIBLE_VCS))
            .arg_from_usage("--hook=[hook]    'Post hook after initialization'")
    }

    fn from_matches<'b: 'a>(m: &'b ArgMatches<'a>) -> NewCommand<'a> {
        NewCommand {
            path: m.value_of("path").unwrap(),
            vcs: m.value_of("vcs")
                .and_then(|s| s.parse().ok())
                .unwrap_or(Vcs::Git),
            hook: m.value_of("hook").and_then(|s| shlex::split(s)),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?;

        let path: Cow<Path> = match self.path.parse::<Query>() {
            Ok(query) => workspace.resolve_query(&query)?.into(),
            Err(_) => Path::new(self.path).into(),
        };

        // init
        let repo = workspace.create_empty_repository(&path, self.vcs)?;

        // hook
        match (repo, self.hook) {
            (Some(ref repo), Some(ref hook)) if hook.len() >= 1 => {
                workspace.print(format_args!("[info] Running post hook command...\n"));
                repo.run_command(&hook[0], &hook[1..])?;
            }
            _ => {}
        }

        workspace.save_cache()?;
        Ok(())
    }
}

/// Subcommand `clone`
pub struct CloneCommand<'a> {
    query: Query,
    dest: Option<PathBuf>,
    root: Option<&'a Path>,
    ssh: bool,
    args: Vec<&'a str>,
    vcs: Vcs,
}

impl<'a> CloneCommand<'a> {
    fn app<'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Clone remote repositories, and then add it under management")
            .arg_from_usage("<query>          'an URL or a string to determine the URL of remote repository'")
            .arg_from_usage("[dest]           'Destination directory of cloned repository'")
            .arg_from_usage("[args]...        'Supplemental arguments for VCS command'")
            .arg_from_usage("--root=[root]    'Path to determine the destination directory of cloned repository'")
            .arg_from_usage("-s, --ssh        'Use SSH protocol'")
            .arg(Arg::from_usage("--vcs=[vcs] 'Used Version Control System'").possible_values(POSSIBLE_VCS))
    }

    fn from_matches<'b: 'a>(m: &'b ArgMatches<'a>) -> CloneCommand<'a> {
        CloneCommand {
            query: m.value_of("query").and_then(|s| s.parse().ok()).unwrap(),
            dest: m.value_of("dest").map(PathBuf::from),
            root: m.value_of("root").map(Path::new),
            ssh: m.is_present("ssh"),
            args: m.values_of("args").map(|s| s.collect()).unwrap_or_default(),
            vcs: m.value_of("vcs")
                .and_then(|s| s.parse().ok())
                .unwrap_or(Vcs::Git),
        }
    }

    fn run(self) -> Result<()> {
        let mut workspace = Workspace::new()?.root_dir(self.root);

        let url = self.query.to_url(self.ssh)?;
        let dest = match self.dest {
            Some(dest) => dest,
            None => workspace.resolve_query(&self.query)?,
        };
        workspace.clone_repository(&url, &dest, self.vcs, &self.args[..])?;

        workspace.save_cache()?;
        Ok(())
    }
}

#[derive(Debug)]
enum ListFormat {
    Name,
    FullPath,
}

impl ::std::str::FromStr for ListFormat {
    type Err = ();
    fn from_str(s: &str) -> ::std::result::Result<Self, Self::Err> {
        match s {
            "name" => Ok(ListFormat::Name),
            "fullpath" => Ok(ListFormat::FullPath),
            _ => Err(()),
        }
    }
}

/// Subcommand `list`
pub struct ListCommand {
    format: ListFormat,
}

impl ListCommand {
    fn app<'a, 'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("List local repositories managed by rhq")
            .arg(clap::Arg::from_usage("--format=[format] 'List format'").possible_values(&["name", "fullpath"]))
    }

    fn from_matches(m: &ArgMatches) -> ListCommand {
        ListCommand {
            format: m.value_of("format")
                .and_then(|s| s.parse().ok())
                .unwrap_or(ListFormat::FullPath),
        }
    }

    fn run(self) -> Result<()> {
        let workspace = Workspace::new()?;
        workspace.for_each_repo(|repo| {
            match self.format {
                ListFormat::Name => println!("{}", repo.name()),
                ListFormat::FullPath => println!("{}", repo.path_string()),
            }
            Ok(())
        })
    }
}

/// Subcommand `foreach`
pub struct ForeachCommand<'a> {
    command: &'a str,
    args: Vec<&'a str>,
    dry_run: bool,
}

impl<'a> ForeachCommand<'a> {
    fn app<'b: 'a>(app: App<'a, 'b>) -> clap::App<'a, 'b> {
        app.about("Execute command into each repositories")
            .arg_from_usage("<command>       'Command name'")
            .arg_from_usage("[args]...       'Arguments of command'")
            .arg_from_usage("-n, --dry-run   'Do not actually execute command'")
    }

    fn from_matches<'b: 'a>(m: &'b ArgMatches<'a>) -> ForeachCommand<'a> {
        ForeachCommand {
            command: m.value_of("command").unwrap(),
            args: m.values_of("args").map(|s| s.collect()).unwrap_or_default(),
            dry_run: m.is_present("dry-run"),
        }
    }

    fn run(self) -> Result<()> {
        let workspace = Workspace::new()?;
        workspace.for_each_repo(|repo| {
            if self.dry_run {
                workspace.print(format_args!(
                    "+ {} {}",
                    self.command,
                    util::join_str(&self.args[..])
                ));
            } else {
                repo.run_command(self.command, &self.args[..])?;
            }
            Ok(())
        })
    }
}

pub struct CompletionCommand<'a> {
    shell: clap::Shell,
    out_file: Option<&'a Path>,
}

impl<'a> CompletionCommand<'a> {
    fn app<'b: 'a>(app: App<'a, 'b>) -> App<'a, 'b> {
        app.about("Generate completion scripts for your shell")
            .setting(AppSettings::ArgRequiredElseHelp)
            .arg(
                clap::Arg::with_name("shell")
                    .help("target shell")
                    .possible_values(&["bash", "zsh", "fish", "powershell"])
                    .required(true),
            )
            .arg(clap::Arg::from_usage("[out-file]  'path to output script'"))
    }

    fn from_matches<'b: 'a>(m: &'b ArgMatches<'a>) -> CompletionCommand<'a> {
        CompletionCommand {
            shell: m.value_of("shell").and_then(|s| s.parse().ok()).unwrap(),
            out_file: m.value_of("out-file").map(Path::new),
        }
    }

    fn run(self) -> Result<()> {
        if let Some(path) = self.out_file {
            let mut file = ::std::fs::OpenOptions::new()
                .write(true)
                .create(true)
                .append(false)
                .open(path)
                .unwrap();
            app().gen_completions_to(env!("CARGO_PKG_NAME"), self.shell, &mut file);
        } else {
            app().gen_completions_to(env!("CARGO_PKG_NAME"), self.shell, &mut ::std::io::stdout());
        }
        Ok(())
    }
}
<|endoftext|>