<commit_before>//! gRPC module for Rust.

use backend::{Initializer, PackageUtils};
use core::errors::Result;
use core::{self, Loc};
use flavored::{RpEndpointHttp1, RpPackage, RpVersionedPackage, RustEndpoint};
use genco::rust::{imported, local};
use genco::{Cons, IntoTokens, Quoted, Rust, Tokens};
use std::rc::Rc;
use utils::Comments;
use {Options, Root, RootCodegen, RustFileSpec, Service, ServiceCodegen, SCOPE_SEP};

pub struct Module {}

impl Module {
    pub fn new() -> Module {
        Module {}
    }
}

impl Initializer for Module {
    type Options = Options;

    fn initialize(&self, options: &mut Options) -> Result<()> {
        let package_utils = options.package_utils.clone();

        let utils_package = RpPackage::new(vec!["reproto".to_string()]);
        let utils_package = RpVersionedPackage::new(utils_package, None);

        let utils_pkg = Rc::new(package_utils.package(&utils_package).join(SCOPE_SEP));
        let result = imported(utils_pkg.clone(), "Result");
        let path_encode = imported(utils_pkg.clone(), "PathEncode");

        options
            .service
            .push(Box::new(ReqwestService::new(result, path_encode)));

        options
            .root
            .push(Box::new(ReqwestUtils::new(utils_package)));

        Ok(())
    }
}

struct ReqwestUtils {
    utils_package: RpVersionedPackage,
}

impl ReqwestUtils {
    pub fn new(utils_package: RpVersionedPackage) -> Self {
        Self { utils_package }
    }

    fn reproto<'el>(&self) -> Result<RustFileSpec<'el>> {
        let mut f = RustFileSpec::default();

        let mut errors = Vec::new();
        errors.push((imported("reqwest", "Error"), "ReqwestError"));
        errors.push((imported("reqwest", "UrlError"), "UrlError"));
        errors.push((imported("std::fmt", "Error"), "FormatError"));

        f.0.push({
            let mut t = Tokens::new();

            push!(t, "pub enum Error {");

            for &(ref ty, ref variant) in &errors {
                nested!(t, variant.clone(), "(", ty.clone(), "),");
            }

            push!(t, "}");

            t
        });

        f.0.push({
            let mut t = Tokens::new();
            let result = imported("std::result", "Result");

            push!(t, "pub type Result<T> = ", result, "<T, Error>;");

            t
        });

        for &(ref ty, ref variant) in &errors {
            f.0.push({
                let mut t = Tokens::new();

                push!(t, "impl From<", ty.clone(), "> for Error {");

                t.nested({
                    let mut t = Tokens::new();

                    push!(t, "fn from(value: ", ty.clone(), ") -> Self {");
                    nested!(t, "Error::", variant.clone(), "(value)");
                    push!(t, "}");

                    t
                });

                push!(t, "}");

                t
            });
        }

        f.0.push({
            let mut t = Tokens::new();

            let display = imported("std::fmt", "Display");
            let fmt = imported("std::fmt", "Formatter");
            let result = imported("std::fmt", "Result");
            let encode = imported("reqwest::header::parsing", "http_percent_encode");

            push!(t, "pub struct PathEncode<T>(pub T);");

            t.push({
                let mut t = Tokens::new();

                push!(t, "impl<T> ", display, " for PathEncode<T>");
                push!(t, "where");
                nested!(t, "T: ", display);
                push!(t, "{");

                t.nested({
                    let mut t = Tokens::new();

                    push!(t, "fn fmt(&self, fmt: &mut ", fmt, ") -> ", result, " {");
                    nested!(t, encode, "(fmt, self.0.to_string().as_bytes())");
                    push!(t, "}");

                    t
                });

                push!(t, "}");

                t
            });

            t.join_line_spacing()
        });

        Ok(f)
    }
}

impl RootCodegen for ReqwestUtils {
    fn generate(&self, root: Root) -> Result<()> {
        let Root { files, .. } = root;
        let package = self.utils_package.clone();
        files.insert(package, self.reproto()?);
        Ok(())
    }
}

struct ReqwestService {
    result: Rust<'static>,
    path_encode: Rust<'static>,
    client: Rust<'static>,
}

impl ReqwestService {
    pub fn new(result: Rust<'static>, path_encode: Rust<'static>) -> Self {
        Self {
            result,
            path_encode,
            client: imported("reqwest", "Client"),
        }
    }
}

impl ServiceCodegen for ReqwestService {
    fn generate(&self, service: Service) -> Result<()> {
        let Service {
            body,
            container,
            name,
            attributes,
            ..
        } = service;

        let name = Cons::from(format!("{}_Reqwest", name));
        let url_ty = imported("reqwest", "Url");

        container.push({
            let mut t = Tokens::new();

            t.push_unless_empty(attributes.clone());
            push!(t, "pub struct ", name, "{");
            nested!(t, "client: ", self.client, ",");
            nested!(t, "url: ", url_ty, ",");
            push!(t, "}");

            t
        });

        container.push({
            let mut t = Tokens::new();

            push!(t, "impl ", name, " {");

            t.push_unless_empty({
                let mut t = Tokens::new();

                // constructor.
                t.nested({
                    let option = local("Option");

                    let mut t = Tokens::new();

                    let mut args = Tokens::new();
                    args.append(toks!["client: ", self.client.clone()]);

                    match body.http.url {
                        Some(_) => {
                            args.append(toks![
                                "url: ",
                                option.with_arguments(vec![url_ty.clone()]),
                            ]);
                        }
                        None => {
                            args.append(toks!["url: ", url_ty.clone()]);
                        }
                    };

                    let s = self.result.with_arguments(vec![local("Self")]);

                    push!(t, "pub fn new(", args.join(", "), ") -> ", s, " {");

                    t.nested({
                        let mut t = Tokens::new();

                        t.push_into(|t| match body.http.url {
                            Some(ref url) => {
                                let url = Loc::value(url).clone().quoted();

                                push!(t, "let url = match url {");
                                nested!(t, "Some(url) => url,");
                                nested!(t, "None => ", url_ty, "::parse(", url, ")?,");
                                push!(t, "};");
                            }
                            None => {
                                push!(t, "let url = Some(url);");
                            }
                        });

                        t.push_into(|t| {
                            push!(t, "Ok(Self {");
                            nested!(t, "client,");
                            nested!(t, "url,");
                            push!(t, "})");
                        });

                        t.join_line_spacing()
                    });

                    push!(t, "}");

                    t
                });

                // endpoint methods.
                for e in &body.endpoints {
                    let http = match e.http1.as_ref() {
                        Some(http) => http,
                        None => continue,
                    };

                    t.nested({
                        let mut t = Tokens::new();

                        t.push_unless_empty(Comments(&e.comment));
                        t.push(Endpoint {
                            result: &self.result,
                            path_encode: &self.path_encode,
                            e,
                            http,
                        });

                        t
                    });
                }

                t.join_line_spacing()
            });

            push!(t, "}");

            t
        });

        Ok(())
    }
}

struct Endpoint<'a, 'el: 'a> {
    result: &'a Rust<'static>,
    path_encode: &'a Rust<'static>,
    e: &'el RustEndpoint,
    http: &'el RpEndpointHttp1,
}

impl<'a, 'el: 'a> IntoTokens<'el, Rust<'el>> for Endpoint<'a, 'el> {
    fn into_tokens(self) -> Tokens<'el, Rust<'el>> {
        let Endpoint {
            result,
            path_encode,
            e,
            http,
        } = self;

        let mut t = Tokens::new();

        // import trait
        t.register(imported("std::fmt", "Write").qualified());

        let mut args = Tokens::new();

        args.append("&self");

        for a in &e.arguments {
            args.append({
                let mut t = Tokens::new();

                t.append(a.safe_ident());
                t.append(": ");
                t.append(a.channel.ty());

                t
            });
        }

        let args = args.join(", ");

        let res = if let Some(ref res) = http.response {
            toks![result.clone(), "<", res, ">"]
        } else {
            toks![result.clone(), "<()>"]
        };

        push!(t, "pub fn ", e.safe_ident(), "(", args, ") -> ", res, " {");

        t.nested({
            let mut t = Tokens::new();

            let method = match http.method {
                core::RpHttpMethod::Get => "Get",
                core::RpHttpMethod::Post => "Post",
                core::RpHttpMethod::Put => "Put",
                core::RpHttpMethod::Update => "Update",
                core::RpHttpMethod::Delete => "Delete",
                core::RpHttpMethod::Patch => "Patch",
                core::RpHttpMethod::Head => "Head",
            };

            let m = toks![imported("reqwest", "Method"), "::", method];

            if let Some(ref http_path) = e.http.path {
                let mut p = Tokens::new();

                push!(p, "let mut path_ = String::new();");

                for step in &http_path.steps {
                    push!(p, "path_.push_str(", "/".quoted(), ");");

                    for part in &step.parts {
                        match *part {
                            core::RpPathPart::Variable(ref arg) => {
                                let var = toks![path_encode.clone(), "(", arg.safe_ident(), ")"];
                                push!(p, "write!(path_, ", "{}".quoted(), ", ", var, ")?;");
                            }
                            core::RpPathPart::Segment(ref s) => {
                                push!(p, "path_.push_str(", s.as_str().quoted(), ");");
                            }
                        }
                    }
                }

                t.push(p);
                push!(t, "let url_ = self.url.join(&path_)?;");
            } else {
                push!(t, "let url_ = self.url.clone();");
            }

            let req = toks!["self.client.request(", m, ", url_)"];

            push!(t, "let mut req_ = ", req, ";");

            if let Some(ref req) = e.request {
                push!(t, "req_.json(&", req.safe_ident(), ");");
            }

            if e.response.is_some() {
                push!(t, "let mut res_ = req_.send()?;");
                push!(t, "let body_ = res_.json()?;");
                push!(t, "Ok(body_)");
            } else {
                push!(t, "req_.send()?;");
                push!(t, "Ok(())");
            }

            t.join_line_spacing()
        });

        push!(t, "}");

        t
    }
}
<commit_msg>rust: reqwest: move Constructor into its own type<commit_after>//! gRPC module for Rust.

use backend::{Initializer, PackageUtils};
use core::errors::Result;
use core::{self, Loc};
use flavored::{RpEndpointHttp1, RpPackage, RpServiceBody, RpVersionedPackage, RustEndpoint};
use genco::rust::{imported, local};
use genco::{Cons, IntoTokens, Quoted, Rust, Tokens};
use std::rc::Rc;
use utils::Comments;
use {Options, Root, RootCodegen, RustFileSpec, Service, ServiceCodegen, SCOPE_SEP};

pub struct Module {}

impl Module {
    pub fn new() -> Module {
        Module {}
    }
}

impl Initializer for Module {
    type Options = Options;

    fn initialize(&self, options: &mut Options) -> Result<()> {
        let package_utils = options.package_utils.clone();

        let utils_package = RpPackage::new(vec!["reproto".to_string()]);
        let utils_package = RpVersionedPackage::new(utils_package, None);

        let utils_pkg = Rc::new(package_utils.package(&utils_package).join(SCOPE_SEP));
        let result = imported(utils_pkg.clone(), "Result");
        let path_encode = imported(utils_pkg.clone(), "PathEncode");

        options
            .service
            .push(Box::new(ReqwestService::new(result, path_encode)));

        options
            .root
            .push(Box::new(ReqwestUtils::new(utils_package)));

        Ok(())
    }
}

struct ReqwestUtils {
    utils_package: RpVersionedPackage,
}

impl ReqwestUtils {
    pub fn new(utils_package: RpVersionedPackage) -> Self {
        Self { utils_package }
    }

    fn reproto<'el>(&self) -> Result<RustFileSpec<'el>> {
        let mut f = RustFileSpec::default();

        let mut errors = Vec::new();
        errors.push((imported("reqwest", "Error"), "ReqwestError"));
        errors.push((imported("reqwest", "UrlError"), "UrlError"));
        errors.push((imported("std::fmt", "Error"), "FormatError"));

        f.0.push({
            let mut t = Tokens::new();

            push!(t, "pub enum Error {");

            for &(ref ty, ref variant) in &errors {
                nested!(t, variant.clone(), "(", ty.clone(), "),");
            }

            push!(t, "}");

            t
        });

        f.0.push({
            let mut t = Tokens::new();
            let result = imported("std::result", "Result");

            push!(t, "pub type Result<T> = ", result, "<T, Error>;");

            t
        });

        for &(ref ty, ref variant) in &errors {
            f.0.push({
                let mut t = Tokens::new();

                push!(t, "impl From<", ty.clone(), "> for Error {");

                t.nested({
                    let mut t = Tokens::new();

                    push!(t, "fn from(value: ", ty.clone(), ") -> Self {");
                    nested!(t, "Error::", variant.clone(), "(value)");
                    push!(t, "}");

                    t
                });

                push!(t, "}");

                t
            });
        }

        f.0.push({
            let mut t = Tokens::new();

            let display = imported("std::fmt", "Display");
            let fmt = imported("std::fmt", "Formatter");
            let result = imported("std::fmt", "Result");
            let encode = imported("reqwest::header::parsing", "http_percent_encode");

            push!(t, "pub struct PathEncode<T>(pub T);");

            t.push({
                let mut t = Tokens::new();

                push!(t, "impl<T> ", display, " for PathEncode<T>");
                push!(t, "where");
                nested!(t, "T: ", display);
                push!(t, "{");

                t.nested({
                    let mut t = Tokens::new();

                    push!(t, "fn fmt(&self, fmt: &mut ", fmt, ") -> ", result, " {");
                    nested!(t, encode, "(fmt, self.0.to_string().as_bytes())");
                    push!(t, "}");

                    t
                });

                push!(t, "}");

                t
            });

            t.join_line_spacing()
        });

        Ok(f)
    }
}

impl RootCodegen for ReqwestUtils {
    fn generate(&self, root: Root) -> Result<()> {
        let Root { files, .. } = root;
        let package = self.utils_package.clone();
        files.insert(package, self.reproto()?);
        Ok(())
    }
}

struct ReqwestService {
    result: Rust<'static>,
    path_encode: Rust<'static>,
    client: Rust<'static>,
}

impl ReqwestService {
    pub fn new(result: Rust<'static>, path_encode: Rust<'static>) -> Self {
        Self {
            result,
            path_encode,
            client: imported("reqwest", "Client"),
        }
    }
}

impl ServiceCodegen for ReqwestService {
    fn generate(&self, service: Service) -> Result<()> {
        let Service {
            body,
            container,
            name,
            attributes,
            ..
        } = service;

        let name = Cons::from(format!("{}_Reqwest", name));
        let url_ty = imported("reqwest", "Url");

        container.push({
            let mut t = Tokens::new();

            t.push_unless_empty(attributes.clone());
            push!(t, "pub struct ", name, "{");
            nested!(t, "client: ", self.client, ",");
            nested!(t, "url: ", url_ty, ",");
            push!(t, "}");

            t
        });

        container.push({
            let mut t = Tokens::new();

            push!(t, "impl ", name, " {");

            t.push_unless_empty({
                let mut t = Tokens::new();

                // constructor.
                t.nested(Constructor {
                    result: &self.result,
                    client: &self.client,
                    body,
                    url_ty: &url_ty,
                });

                // endpoint methods.
                for e in &body.endpoints {
                    let http = match e.http1.as_ref() {
                        Some(http) => http,
                        None => continue,
                    };

                    t.nested({
                        let mut t = Tokens::new();

                        t.push_unless_empty(Comments(&e.comment));
                        t.push(Endpoint {
                            result: &self.result,
                            path_encode: &self.path_encode,
                            e,
                            http,
                        });

                        t
                    });
                }

                t.join_line_spacing()
            });

            push!(t, "}");

            t
        });

        Ok(())
    }
}

/// Builds a constructor for the service struct.
struct Constructor<'a, 'el: 'a> {
    body: &'el RpServiceBody,
    result: &'a Rust<'static>,
    client: &'a Rust<'static>,
    url_ty: &'a Rust<'static>,
}

impl<'a, 'el: 'a> IntoTokens<'el, Rust<'el>> for Constructor<'a, 'el> {
    fn into_tokens(self) -> Tokens<'el, Rust<'el>> {
        let Constructor {
            body,
            result,
            client,
            url_ty,
            ..
        } = self;

        let option = local("Option");

        let mut t = Tokens::new();

        let mut args = Tokens::new();
        args.append(toks!["client: ", client.clone()]);

        match body.http.url {
            Some(_) => {
                args.append(toks!["url: ", option.with_arguments(vec![url_ty.clone()]),]);
            }
            None => {
                args.append(toks!["url: ", url_ty.clone()]);
            }
        };

        let s = result.with_arguments(vec![local("Self")]);

        push!(t, "pub fn new(", args.join(", "), ") -> ", s, " {");

        t.nested({
            let mut t = Tokens::new();

            t.push_into(|t| match body.http.url {
                Some(ref url) => {
                    let url = Loc::value(url).clone().quoted();

                    push!(t, "let url = match url {");
                    nested!(t, "Some(url) => url,");
                    nested!(t, "None => ", url_ty.clone(), "::parse(", url, ")?,");
                    push!(t, "};");
                }
                None => {
                    push!(t, "let url = Some(url);");
                }
            });

            t.push_into(|t| {
                push!(t, "Ok(Self {");
                nested!(t, "client,");
                nested!(t, "url,");
                push!(t, "})");
            });

            t.join_line_spacing()
        });

        push!(t, "}");

        t
    }
}

/// Build an endpoint method for the service struct.
struct Endpoint<'a, 'el: 'a> {
    result: &'a Rust<'static>,
    path_encode: &'a Rust<'static>,
    e: &'el RustEndpoint,
    http: &'el RpEndpointHttp1,
}

impl<'a, 'el: 'a> IntoTokens<'el, Rust<'el>> for Endpoint<'a, 'el> {
    fn into_tokens(self) -> Tokens<'el, Rust<'el>> {
        let Endpoint {
            result,
            path_encode,
            e,
            http,
        } = self;

        let mut t = Tokens::new();

        // import trait
        t.register(imported("std::fmt", "Write").qualified());

        let mut args = Tokens::new();

        args.append("&self");

        for a in &e.arguments {
            args.append({
                let mut t = Tokens::new();

                t.append(a.safe_ident());
                t.append(": ");
                t.append(a.channel.ty());

                t
            });
        }

        let args = args.join(", ");

        let res = if let Some(ref res) = http.response {
            toks![result.clone(), "<", res, ">"]
        } else {
            toks![result.clone(), "<()>"]
        };

        push!(t, "pub fn ", e.safe_ident(), "(", args, ") -> ", res, " {");

        t.nested({
            let mut t = Tokens::new();

            let method = match http.method {
                core::RpHttpMethod::Get => "Get",
                core::RpHttpMethod::Post => "Post",
                core::RpHttpMethod::Put => "Put",
                core::RpHttpMethod::Update => "Update",
                core::RpHttpMethod::Delete => "Delete",
                core::RpHttpMethod::Patch => "Patch",
                core::RpHttpMethod::Head => "Head",
            };

            let m = toks![imported("reqwest", "Method"), "::", method];

            if let Some(ref http_path) = e.http.path {
                let mut p = Tokens::new();

                push!(p, "let mut path_ = String::new();");

                for step in &http_path.steps {
                    push!(p, "path_.push_str(", "/".quoted(), ");");

                    for part in &step.parts {
                        match *part {
                            core::RpPathPart::Variable(ref arg) => {
                                let var = toks![path_encode.clone(), "(", arg.safe_ident(), ")"];
                                push!(p, "write!(path_, ", "{}".quoted(), ", ", var, ")?;");
                            }
                            core::RpPathPart::Segment(ref s) => {
                                push!(p, "path_.push_str(", s.as_str().quoted(), ");");
                            }
                        }
                    }
                }

                t.push(p);
                push!(t, "let url_ = self.url.join(&path_)?;");
            } else {
                push!(t, "let url_ = self.url.clone();");
            }

            let req = toks!["self.client.request(", m, ", url_)"];

            push!(t, "let mut req_ = ", req, ";");

            if let Some(ref req) = e.request {
                push!(t, "req_.json(&", req.safe_ident(), ");");
            }

            if e.response.is_some() {
                push!(t, "let mut res_ = req_.send()?;");
                push!(t, "let body_ = res_.json()?;");
                push!(t, "Ok(body_)");
            } else {
                push!(t, "req_.send()?;");
                push!(t, "Ok(())");
            }

            t.join_line_spacing()
        });

        push!(t, "}");

        t
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2015-2016 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Implementation of the Bulk API

use hyper::status::StatusCode;

use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde::de::{Error, MapVisitor, Visitor};
use serde_json;

use ::{Client, EsResponse};
use ::do_req;
use ::error::EsError;
use ::json::{FieldBased, NoOuter, ShouldSkip};
use ::units::Duration;

use super::ShardCountResult;
use super::common::{Options, OptionVal, VersionType};

pub enum ActionType {
    Index,
    Create,
    Delete,
    /// WARNING - currently un-implemented
    Update
}

impl Serialize for ActionType {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
        where S: Serializer {

        self.to_string().serialize(serializer)
    }
}

impl ToString for ActionType {
    fn to_string(&self) -> String {
        match *self {
            ActionType::Index => "index",
            ActionType::Create => "create",
            ActionType::Delete => "delete",
            ActionType::Update => "update"
        }.to_owned()
    }
}

#[derive(Default, Serialize)]
pub struct ActionOptions {
    #[serde(rename="_index", skip_serializing_if="ShouldSkip::should_skip")]
    index:             Option<String>,
    #[serde(rename="_type", skip_serializing_if="ShouldSkip::should_skip")]
    doc_type:          Option<String>,
    #[serde(rename="_id", skip_serializing_if="ShouldSkip::should_skip")]
    id:                Option<String>,
    #[serde(rename="_version", skip_serializing_if="ShouldSkip::should_skip")]
    version:           Option<u64>,
    #[serde(rename="_version_type", skip_serializing_if="ShouldSkip::should_skip")]
    version_type:      Option<VersionType>,
    #[serde(rename="_routing", skip_serializing_if="ShouldSkip::should_skip")]
    routing:           Option<String>,
    #[serde(rename="_parent", skip_serializing_if="ShouldSkip::should_skip")]
    parent:            Option<String>,
    #[serde(rename="_timestamp", skip_serializing_if="ShouldSkip::should_skip")]
    timestamp:         Option<String>,
    #[serde(rename="_ttl", skip_serializing_if="ShouldSkip::should_skip")]
    ttl:               Option<Duration>,
    #[serde(rename="_retry_on_conflict", skip_serializing_if="ShouldSkip::should_skip")]
    retry_on_conflict: Option<u64>,
}

#[derive(Serialize)]
pub struct Action<X>(FieldBased<ActionType, ActionOptions, NoOuter>, Option<X>);

impl<S> Action<S>
    where S: Serialize {
    /// An index action.
    ///
    /// Takes the document to be indexed, other parameters can be set as
    /// optional on the `Action` struct returned.
    pub fn index(document: S) -> Self {
        Action(FieldBased::new(ActionType::Index,
                               Default::default(),
                               NoOuter),
               Some(document))
    }

    /// Create action
    pub fn create(document: S) -> Self {
        Action(FieldBased::new(ActionType::Create,
                               Default::default(),
                               NoOuter),
               Some(document))
    }

    /// Add the serialized version of this action to the bulk `String`.
    fn add(&self, actstr: &mut String) -> Result<(), EsError> {
        let command_str = try!(serde_json::to_string(&self.0));

        actstr.push_str(&command_str);
        actstr.push_str("\n");

        match self.1 {
            Some(ref source) => {
                let payload_str = try!(serde_json::to_string(source));
                actstr.push_str(&payload_str);
                actstr.push_str("\n");
            },
            None             => ()
        }
        Ok(())
    }
}

impl<S> Action<S> {
    /// Delete a document based on ID.
    ///
    /// # Example
    ///
    /// ```
    /// use rs_es::operations::bulk::Action;
    ///
    /// let delete_action:Action<()> = Action::delete("doc_id");
    /// let delete_with_index:Action<()> = Action::delete("doc_id").with_index("index_name");
    /// ```
    pub fn delete<A: Into<String>>(id: A) -> Self {
        Action(FieldBased::new(ActionType::Delete,
                               ActionOptions {
                                   id: Some(id.into()),
                                   ..Default::default()
                               },
                               NoOuter),
               None)
    }

    // TODO - implement update

    add_inner_field!(with_index, index, String);
    add_inner_field!(with_doc_type, doc_type, String);
    add_inner_field!(with_id, id, String);
    add_inner_field!(with_version, version, u64);
    add_inner_field!(with_version_type, version_type, VersionType);
    add_inner_field!(with_routing, routing, String);
    add_inner_field!(with_parent, parent, String);
    add_inner_field!(with_timestamp, timestamp, String);
    add_inner_field!(with_ttl, ttl, Duration);
    add_inner_field!(with_retry_on_conflict, retry_on_conflict, u64);
}

pub struct BulkOperation<'a, 'b, S: 'b> {
    client:   &'a mut Client,
    index:    Option<&'b str>,
    doc_type: Option<&'b str>,
    actions:  &'b [Action<S>],
    options:  Options<'b>
}

impl<'a, 'b, S> BulkOperation<'a, 'b, S>
    where S: Serialize {

    pub fn new(client: &'a mut Client, actions: &'b [Action<S>]) -> Self {
        BulkOperation {
            client:   client,
            index:    None,
            doc_type: None,
            actions:  actions,
            options:  Options::new()
        }
    }

    pub fn with_index(&'b mut self, index: &'b str) -> &'b mut Self {
        self.index = Some(index);
        self
    }

    pub fn with_doc_type(&'b mut self, doc_type: &'b str) -> &'b mut Self {
        self.doc_type = Some(doc_type);
        self
    }

    add_option!(with_consistency, "consistency");
    add_option!(with_refresh, "refresh");

    fn format_url(&self) -> String {
        let mut url = String::new();
        url.push_str("/");
        match self.index {
            Some(index) => {
                url.push_str(index);
                url.push_str("/");
            },
            None        => ()
        }
        match self.doc_type {
            Some(doc_type) => {
                url.push_str(doc_type);
                url.push_str("/");
            },
            None           => ()
        }
        url.push_str("_bulk");
        url.push_str(&self.options.to_string());
        url
    }

    fn format_actions(&self) -> String {
        let mut actstr = String::new();
        for action in self.actions {
            action.add(&mut actstr).unwrap();
        }
        actstr
    }

    pub fn send(&'b mut self) -> Result<BulkResult, EsError> {
        //
        // This function does not use the standard GET/POST/DELETE functions of
        // the client, as they serve the happy path of JSON-in/JSON-out, this
        // function does send send JSON in.
        //
        // Various parts of the client are reused where it makes sense.
        //
        let full_url = {
            let url = self.format_url();
            self.client.full_url(&url)
        };
        let body = self.format_actions();
        debug!("Sending: {}", body);
        // Doesn't use the standard macros as it's not standard JSON
        let result = try!(self.client.http_client
                          .post(&full_url)
                          .body(&body)
                          .send());

        let response = try!(do_req(result));

        match response.status_code() {
            &StatusCode::Ok => Ok(try!(response.read_response())),
            _              => Err(EsError::EsError(format!("Unexpected status: {}", response.status_code())))
        }
    }
}

impl Client {
    /// Bulk
    ///
    /// See: https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html
    pub fn bulk<'a, 'b, S>(&'a mut self,
                           actions: &'b [Action<S>]) -> BulkOperation<'a, 'b, S>
        where S: Serialize {

        BulkOperation::new(self, actions)
    }
}

/// The result of specific actions
pub struct ActionResult {
    pub action: ActionType,
    pub inner: ActionResultInner
}

impl Deserialize for ActionResult {
    fn deserialize<D>(deserializer: &mut D) -> Result<ActionResult, D::Error>
        where D: Deserializer {

        struct ActionResultVisitor;

        impl Visitor for ActionResultVisitor {
            type Value = ActionResult;

            fn visit_map<V>(&mut self, mut visitor: V) -> Result<ActionResult, V::Error>
                where V: MapVisitor {

                let visited:Option<(String, ActionResultInner)> = try!(visitor.visit());
                let (key, value) = match visited {
                    Some((key, value)) => (key, value),
                    None               => return Err(V::Error::custom("expecting at least one field"))
                };
                try!(visitor.end());

                let result = ActionResult {
                    action: match key.as_ref() {
                        "index" => ActionType::Index,
                        "create" => ActionType::Create,
                        "delete" => ActionType::Delete,
                        "update" => ActionType::Update,
                        _ => {
                            return Err(V::Error::custom(format!("Unrecognised key: {}", key)))
                        }
                    },
                    inner:  value
                };

                Ok(result)
            }
        }

        deserializer.deserialize(ActionResultVisitor)
    }
}

#[derive(Debug, Deserialize)]
pub struct ActionResultInner {
    #[serde(rename="_index")]
    pub index:    String,
    #[serde(rename="_type")]
    pub doc_type: String,
    #[serde(rename="_version")]
    pub version:  u64,
    pub status:   u64,
    #[serde(rename="_shards")]
    pub shards:   ShardCountResult,
    pub found:    Option<bool>
}

/// The result of a bulk operation
#[derive(Deserialize)]
pub struct BulkResult {
    pub errors: bool,
    pub items:  Vec<ActionResult>,
    pub took:   u64
}

#[cfg(test)]
pub mod tests {
    use ::tests::{clean_db, TestDocument, make_client};

    use super::Action;

    #[test]
    fn test_bulk() {
        let index_name = "test_bulk";
        let mut client = make_client();

        clean_db(&mut client, index_name);

        let actions:Vec<Action<TestDocument>> = (1..10).map(|i| {
            let doc = TestDocument::new().with_str_field("bulk_doc").with_int_field(i);
            Action::index(doc)
        }).collect();

        let result = client.bulk(&actions)
            .with_index(index_name)
            .with_doc_type("bulk_type")
            .send()
            .unwrap();

        assert_eq!(false, result.errors);
        assert_eq!(9, result.items.len());
    }
}
<commit_msg>Send headers in bulk operations<commit_after>/*
 * Copyright 2015-2016 Ben Ashford
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

//! Implementation of the Bulk API

use hyper::status::StatusCode;

use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde::de::{Error, MapVisitor, Visitor};
use serde_json;

use ::{Client, EsResponse};
use ::do_req;
use ::error::EsError;
use ::json::{FieldBased, NoOuter, ShouldSkip};
use ::units::Duration;

use super::ShardCountResult;
use super::common::{Options, OptionVal, VersionType};

pub enum ActionType {
    Index,
    Create,
    Delete,
    /// WARNING - currently un-implemented
    Update
}

impl Serialize for ActionType {
    fn serialize<S>(&self, serializer: &mut S) -> Result<(), S::Error>
        where S: Serializer {

        self.to_string().serialize(serializer)
    }
}

impl ToString for ActionType {
    fn to_string(&self) -> String {
        match *self {
            ActionType::Index => "index",
            ActionType::Create => "create",
            ActionType::Delete => "delete",
            ActionType::Update => "update"
        }.to_owned()
    }
}

#[derive(Default, Serialize)]
pub struct ActionOptions {
    #[serde(rename="_index", skip_serializing_if="ShouldSkip::should_skip")]
    index:             Option<String>,
    #[serde(rename="_type", skip_serializing_if="ShouldSkip::should_skip")]
    doc_type:          Option<String>,
    #[serde(rename="_id", skip_serializing_if="ShouldSkip::should_skip")]
    id:                Option<String>,
    #[serde(rename="_version", skip_serializing_if="ShouldSkip::should_skip")]
    version:           Option<u64>,
    #[serde(rename="_version_type", skip_serializing_if="ShouldSkip::should_skip")]
    version_type:      Option<VersionType>,
    #[serde(rename="_routing", skip_serializing_if="ShouldSkip::should_skip")]
    routing:           Option<String>,
    #[serde(rename="_parent", skip_serializing_if="ShouldSkip::should_skip")]
    parent:            Option<String>,
    #[serde(rename="_timestamp", skip_serializing_if="ShouldSkip::should_skip")]
    timestamp:         Option<String>,
    #[serde(rename="_ttl", skip_serializing_if="ShouldSkip::should_skip")]
    ttl:               Option<Duration>,
    #[serde(rename="_retry_on_conflict", skip_serializing_if="ShouldSkip::should_skip")]
    retry_on_conflict: Option<u64>,
}

#[derive(Serialize)]
pub struct Action<X>(FieldBased<ActionType, ActionOptions, NoOuter>, Option<X>);

impl<S> Action<S>
    where S: Serialize {
    /// An index action.
    ///
    /// Takes the document to be indexed, other parameters can be set as
    /// optional on the `Action` struct returned.
    pub fn index(document: S) -> Self {
        Action(FieldBased::new(ActionType::Index,
                               Default::default(),
                               NoOuter),
               Some(document))
    }

    /// Create action
    pub fn create(document: S) -> Self {
        Action(FieldBased::new(ActionType::Create,
                               Default::default(),
                               NoOuter),
               Some(document))
    }

    /// Add the serialized version of this action to the bulk `String`.
    fn add(&self, actstr: &mut String) -> Result<(), EsError> {
        let command_str = try!(serde_json::to_string(&self.0));

        actstr.push_str(&command_str);
        actstr.push_str("\n");

        match self.1 {
            Some(ref source) => {
                let payload_str = try!(serde_json::to_string(source));
                actstr.push_str(&payload_str);
                actstr.push_str("\n");
            },
            None             => ()
        }
        Ok(())
    }
}

impl<S> Action<S> {
    /// Delete a document based on ID.
    ///
    /// # Example
    ///
    /// ```
    /// use rs_es::operations::bulk::Action;
    ///
    /// let delete_action:Action<()> = Action::delete("doc_id");
    /// let delete_with_index:Action<()> = Action::delete("doc_id").with_index("index_name");
    /// ```
    pub fn delete<A: Into<String>>(id: A) -> Self {
        Action(FieldBased::new(ActionType::Delete,
                               ActionOptions {
                                   id: Some(id.into()),
                                   ..Default::default()
                               },
                               NoOuter),
               None)
    }

    // TODO - implement update

    add_inner_field!(with_index, index, String);
    add_inner_field!(with_doc_type, doc_type, String);
    add_inner_field!(with_id, id, String);
    add_inner_field!(with_version, version, u64);
    add_inner_field!(with_version_type, version_type, VersionType);
    add_inner_field!(with_routing, routing, String);
    add_inner_field!(with_parent, parent, String);
    add_inner_field!(with_timestamp, timestamp, String);
    add_inner_field!(with_ttl, ttl, Duration);
    add_inner_field!(with_retry_on_conflict, retry_on_conflict, u64);
}

pub struct BulkOperation<'a, 'b, S: 'b> {
    client:   &'a mut Client,
    index:    Option<&'b str>,
    doc_type: Option<&'b str>,
    actions:  &'b [Action<S>],
    options:  Options<'b>
}

impl<'a, 'b, S> BulkOperation<'a, 'b, S>
    where S: Serialize {

    pub fn new(client: &'a mut Client, actions: &'b [Action<S>]) -> Self {
        BulkOperation {
            client:   client,
            index:    None,
            doc_type: None,
            actions:  actions,
            options:  Options::new()
        }
    }

    pub fn with_index(&'b mut self, index: &'b str) -> &'b mut Self {
        self.index = Some(index);
        self
    }

    pub fn with_doc_type(&'b mut self, doc_type: &'b str) -> &'b mut Self {
        self.doc_type = Some(doc_type);
        self
    }

    add_option!(with_consistency, "consistency");
    add_option!(with_refresh, "refresh");

    fn format_url(&self) -> String {
        let mut url = String::new();
        url.push_str("/");
        match self.index {
            Some(index) => {
                url.push_str(index);
                url.push_str("/");
            },
            None        => ()
        }
        match self.doc_type {
            Some(doc_type) => {
                url.push_str(doc_type);
                url.push_str("/");
            },
            None           => ()
        }
        url.push_str("_bulk");
        url.push_str(&self.options.to_string());
        url
    }

    fn format_actions(&self) -> String {
        let mut actstr = String::new();
        for action in self.actions {
            action.add(&mut actstr).unwrap();
        }
        actstr
    }

    pub fn send(&'b mut self) -> Result<BulkResult, EsError> {
        //
        // This function does not use the standard GET/POST/DELETE functions of
        // the client, as they serve the happy path of JSON-in/JSON-out, this
        // function does send send JSON in.
        //
        // Various parts of the client are reused where it makes sense.
        //
        let full_url = {
            let url = self.format_url();
            self.client.full_url(&url)
        };
        let body = self.format_actions();
        debug!("Sending: {}", body);
        // Doesn't use the standard macros as it's not standard JSON
        let result = try!(self.client.http_client
                          .post(&full_url)
                          .body(&body)
                          .headers(self.client.headers.clone())
                          .send());

        let response = try!(do_req(result));

        match response.status_code() {
            &StatusCode::Ok => Ok(try!(response.read_response())),
            _              => Err(EsError::EsError(format!("Unexpected status: {}", response.status_code())))
        }
    }
}

impl Client {
    /// Bulk
    ///
    /// See: https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html
    pub fn bulk<'a, 'b, S>(&'a mut self,
                           actions: &'b [Action<S>]) -> BulkOperation<'a, 'b, S>
        where S: Serialize {

        BulkOperation::new(self, actions)
    }
}

/// The result of specific actions
pub struct ActionResult {
    pub action: ActionType,
    pub inner: ActionResultInner
}

impl Deserialize for ActionResult {
    fn deserialize<D>(deserializer: &mut D) -> Result<ActionResult, D::Error>
        where D: Deserializer {

        struct ActionResultVisitor;

        impl Visitor for ActionResultVisitor {
            type Value = ActionResult;

            fn visit_map<V>(&mut self, mut visitor: V) -> Result<ActionResult, V::Error>
                where V: MapVisitor {

                let visited:Option<(String, ActionResultInner)> = try!(visitor.visit());
                let (key, value) = match visited {
                    Some((key, value)) => (key, value),
                    None               => return Err(V::Error::custom("expecting at least one field"))
                };
                try!(visitor.end());

                let result = ActionResult {
                    action: match key.as_ref() {
                        "index" => ActionType::Index,
                        "create" => ActionType::Create,
                        "delete" => ActionType::Delete,
                        "update" => ActionType::Update,
                        _ => {
                            return Err(V::Error::custom(format!("Unrecognised key: {}", key)))
                        }
                    },
                    inner:  value
                };

                Ok(result)
            }
        }

        deserializer.deserialize(ActionResultVisitor)
    }
}

#[derive(Debug, Deserialize)]
pub struct ActionResultInner {
    #[serde(rename="_index")]
    pub index:    String,
    #[serde(rename="_type")]
    pub doc_type: String,
    #[serde(rename="_version")]
    pub version:  u64,
    pub status:   u64,
    #[serde(rename="_shards")]
    pub shards:   ShardCountResult,
    pub found:    Option<bool>
}

/// The result of a bulk operation
#[derive(Deserialize)]
pub struct BulkResult {
    pub errors: bool,
    pub items:  Vec<ActionResult>,
    pub took:   u64
}

#[cfg(test)]
pub mod tests {
    use ::tests::{clean_db, TestDocument, make_client};

    use super::Action;

    #[test]
    fn test_bulk() {
        let index_name = "test_bulk";
        let mut client = make_client();

        clean_db(&mut client, index_name);

        let actions:Vec<Action<TestDocument>> = (1..10).map(|i| {
            let doc = TestDocument::new().with_str_field("bulk_doc").with_int_field(i);
            Action::index(doc)
        }).collect();

        let result = client.bulk(&actions)
            .with_index(index_name)
            .with_doc_type("bulk_type")
            .send()
            .unwrap();

        assert_eq!(false, result.errors);
        assert_eq!(9, result.items.len());
    }
}
<|endoftext|>


<commit_before>use output::Colours;

use options::{flags, Misfire};
use options::parser::MatchedFlags;

/// Under what circumstances we should display coloured, rather than plain,
/// output to the terminal.
///
/// By default, we want to display the colours when stdout can display them.
/// Turning them on when output is going to, say, a pipe, would make programs
/// such as `grep` or `more` not work properly. So the `Automatic` mode does
/// this check and only displays colours when they can be truly appreciated.
#[derive(PartialEq, Debug)]
enum TerminalColours {

    /// Display them even when output isn’t going to a terminal.
    Always,

    /// Display them when output is going to a terminal, but not otherwise.
    Automatic,

    /// Never display them, even when output is going to a terminal.
    Never,
}

impl Default for TerminalColours {
    fn default() -> TerminalColours {
        TerminalColours::Automatic
    }
}

const COLOURS: &[&str] = &["always", "auto", "never"];

impl TerminalColours {

    /// Determine which terminal colour conditions to use.
    fn deduce(matches: &MatchedFlags) -> Result<TerminalColours, Misfire> {

        let word = match matches.get_where(|f| f.matches(&flags::COLOR) || f.matches(&flags::COLOUR))? {
            Some(w) => w,
            None    => return Ok(TerminalColours::default()),
        };

        if word == "always" {
            Ok(TerminalColours::Always)
        }
        else if word == "auto" || word == "automatic" {
            Ok(TerminalColours::Automatic)
        }
        else if word == "never" {
            Ok(TerminalColours::Never)
        }
        else {
            Err(Misfire::bad_argument(&flags::COLOR, word, COLOURS))
        }
    }
}

impl Colours {
    pub fn deduce<TW>(matches: &MatchedFlags, widther: TW) -> Result<Colours, Misfire>
    where TW: Fn() -> Option<usize> {
        use self::TerminalColours::*;

        let tc = TerminalColours::deduce(matches)?;
        if tc == Always || (tc == Automatic && widther().is_some()) {
            let scale = matches.has_where(|f| f.matches(&flags::COLOR_SCALE) || f.matches(&flags::COLOUR_SCALE))?;
            Ok(Colours::colourful(scale.is_some()))
        }
        else {
            Ok(Colours::plain())
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use std::ffi::OsString;
    use options::flags;
    use options::parser::{Flag, Arg};

    use options::test::parse_for_test;
    use options::test::Strictnesses::*;

    pub fn os(input: &'static str) -> OsString {
        let mut os = OsString::new();
        os.push(input);
        os
    }

    static TEST_ARGS: &[&Arg] = &[ &flags::COLOR,       &flags::COLOUR,
                                   &flags::COLOR_SCALE, &flags::COLOUR_SCALE ];

    macro_rules! test {

        ($name:ident: $type:ident <- $inputs:expr; $stricts:expr => $result:expr) => {
            /// Macro that writes a test.
            /// If testing both strictnesses, they’ll both be done in the same function.
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf)) {
                    assert_eq!(result, $result);
                }
            }
        };

        ($name:ident: $type:ident <- $inputs:expr; $stricts:expr => err $result:expr) => {
            /// Special macro for testing Err results.
            /// This is needed because sometimes the Ok type doesn’t implement PartialEq.
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf)) {
                    assert_eq!(result.unwrap_err(), $result);
                }
            }
        };

        ($name:ident: $type:ident <- $inputs:expr; $stricts:expr => like $pat:pat) => {
            /// More general macro for testing against a pattern.
            /// Instead of using PartialEq, this just tests if it matches a pat.
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf)) {
                    println!("Testing {:?}", result);
                    match result {
                        $pat => assert!(true),
                        _    => assert!(false),
                    }
                }
            }
        };

        ($name:ident: $type:ident <- $inputs:expr, $vars:expr; $stricts:expr => err $result:expr) => {
            /// Like above, but with $vars.
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf, &$vars)) {
                    assert_eq!(result.unwrap_err(), $result);
                }
            }
        };

        ($name:ident: $type:ident <- $inputs:expr, $vars:expr; $stricts:expr => like $pat:pat) => {
            /// Like further above, but with $vars.
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf, &$vars)) {
                    println!("Testing {:?}", result);
                    match result {
                        $pat => assert!(true),
                        _    => assert!(false),
                    }
                }
            }
        };
    }

    // Default
    test!(empty:        TerminalColours <- [];                     Both => Ok(TerminalColours::default()));

    // --colour
    test!(u_always:     TerminalColours <- ["--colour=always"];    Both => Ok(TerminalColours::Always));
    test!(u_auto:       TerminalColours <- ["--colour", "auto"];   Both => Ok(TerminalColours::Automatic));
    test!(u_never:      TerminalColours <- ["--colour=never"];     Both => Ok(TerminalColours::Never));

    // --color
    test!(no_u_always:  TerminalColours <- ["--color", "always"];  Both => Ok(TerminalColours::Always));
    test!(no_u_auto:    TerminalColours <- ["--color=auto"];       Both => Ok(TerminalColours::Automatic));
    test!(no_u_never:   TerminalColours <- ["--color", "never"];   Both => Ok(TerminalColours::Never));

    // Errors
    test!(no_u_error:   TerminalColours <- ["--color=upstream"];   Both => err Misfire::bad_argument(&flags::COLOR, &os("upstream"), super::COLOURS));  // the error is for --color
    test!(u_error:      TerminalColours <- ["--colour=lovers"];    Both => err Misfire::bad_argument(&flags::COLOR, &os("lovers"),   super::COLOURS));  // and so is this one!

    // Overriding
    test!(overridden_1: TerminalColours <- ["--colour=auto", "--colour=never"];  Last => Ok(TerminalColours::Never));
    test!(overridden_2: TerminalColours <- ["--color=auto",  "--colour=never"];  Last => Ok(TerminalColours::Never));
    test!(overridden_3: TerminalColours <- ["--colour=auto", "--color=never"];   Last => Ok(TerminalColours::Never));
    test!(overridden_4: TerminalColours <- ["--color=auto",  "--color=never"];   Last => Ok(TerminalColours::Never));

    test!(overridden_5: TerminalColours <- ["--colour=auto", "--colour=never"];  Complain => err Misfire::Duplicate(Flag::Long("colour"), Flag::Long("colour")));
    test!(overridden_6: TerminalColours <- ["--color=auto",  "--colour=never"];  Complain => err Misfire::Duplicate(Flag::Long("color"),  Flag::Long("colour")));
    test!(overridden_7: TerminalColours <- ["--colour=auto", "--color=never"];   Complain => err Misfire::Duplicate(Flag::Long("colour"), Flag::Long("color")));
    test!(overridden_8: TerminalColours <- ["--color=auto",  "--color=never"];   Complain => err Misfire::Duplicate(Flag::Long("color"),  Flag::Long("color")));

    test!(scale_1: Colours <- ["--color=always", "--color-scale", "--colour-scale"], || None;   Last => like Ok(Colours { scale: true,  .. }));
    test!(scale_2: Colours <- ["--color=always", "--color-scale",                 ], || None;   Last => like Ok(Colours { scale: true,  .. }));
    test!(scale_3: Colours <- ["--color=always",                  "--colour-scale"], || None;   Last => like Ok(Colours { scale: true,  .. }));
    test!(scale_4: Colours <- ["--color=always",                                  ], || None;   Last => like Ok(Colours { scale: false, .. }));

    test!(scale_5: Colours <- ["--color=always", "--color-scale", "--colour-scale"], || None;   Complain => err Misfire::Duplicate(Flag::Long("color-scale"),  Flag::Long("colour-scale")));
    test!(scale_6: Colours <- ["--color=always", "--color-scale",                 ], || None;   Complain => like Ok(Colours { scale: true,  .. }));
    test!(scale_7: Colours <- ["--color=always",                  "--colour-scale"], || None;   Complain => like Ok(Colours { scale: true,  .. }));
    test!(scale_8: Colours <- ["--color=always",                                  ], || None;   Complain => like Ok(Colours { scale: false, .. }));
}
<commit_msg>Split the colours tests in two<commit_after>use output::Colours;

use options::{flags, Misfire};
use options::parser::MatchedFlags;

/// Under what circumstances we should display coloured, rather than plain,
/// output to the terminal.
///
/// By default, we want to display the colours when stdout can display them.
/// Turning them on when output is going to, say, a pipe, would make programs
/// such as `grep` or `more` not work properly. So the `Automatic` mode does
/// this check and only displays colours when they can be truly appreciated.
#[derive(PartialEq, Debug)]
enum TerminalColours {

    /// Display them even when output isn’t going to a terminal.
    Always,

    /// Display them when output is going to a terminal, but not otherwise.
    Automatic,

    /// Never display them, even when output is going to a terminal.
    Never,
}

impl Default for TerminalColours {
    fn default() -> TerminalColours {
        TerminalColours::Automatic
    }
}

const COLOURS: &[&str] = &["always", "auto", "never"];

impl TerminalColours {

    /// Determine which terminal colour conditions to use.
    fn deduce(matches: &MatchedFlags) -> Result<TerminalColours, Misfire> {

        let word = match matches.get_where(|f| f.matches(&flags::COLOR) || f.matches(&flags::COLOUR))? {
            Some(w) => w,
            None    => return Ok(TerminalColours::default()),
        };

        if word == "always" {
            Ok(TerminalColours::Always)
        }
        else if word == "auto" || word == "automatic" {
            Ok(TerminalColours::Automatic)
        }
        else if word == "never" {
            Ok(TerminalColours::Never)
        }
        else {
            Err(Misfire::bad_argument(&flags::COLOR, word, COLOURS))
        }
    }
}

impl Colours {
    pub fn deduce<TW>(matches: &MatchedFlags, widther: TW) -> Result<Colours, Misfire>
    where TW: Fn() -> Option<usize> {
        use self::TerminalColours::*;

        let tc = TerminalColours::deduce(matches)?;
        if tc == Always || (tc == Automatic && widther().is_some()) {
            let scale = matches.has_where(|f| f.matches(&flags::COLOR_SCALE) || f.matches(&flags::COLOUR_SCALE))?;
            Ok(Colours::colourful(scale.is_some()))
        }
        else {
            Ok(Colours::plain())
        }
    }
}

#[cfg(test)]
mod terminal_test {
    use super::*;
    use std::ffi::OsString;
    use options::flags;
    use options::parser::{Flag, Arg};

    use options::test::parse_for_test;
    use options::test::Strictnesses::*;

    pub fn os(input: &'static str) -> OsString {
        let mut os = OsString::new();
        os.push(input);
        os
    }

    static TEST_ARGS: &[&Arg] = &[ &flags::COLOR, &flags::COLOUR ];

    macro_rules! test {
        ($name:ident: $type:ident <- $inputs:expr; $stricts:expr => $result:expr) => {
            /// Macro that writes a test.
            /// If testing both strictnesses, they’ll both be done in the same function.
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf)) {
                    assert_eq!(result, $result);
                }
            }
        };

        ($name:ident: $type:ident <- $inputs:expr; $stricts:expr => err $result:expr) => {
            /// Special macro for testing Err results.
            /// This is needed because sometimes the Ok type doesn’t implement PartialEq.
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf)) {
                    assert_eq!(result.unwrap_err(), $result);
                }
            }
        };
    }

    // Default
    test!(empty:        TerminalColours <- [];                     Both => Ok(TerminalColours::default()));

    // --colour
    test!(u_always:     TerminalColours <- ["--colour=always"];    Both => Ok(TerminalColours::Always));
    test!(u_auto:       TerminalColours <- ["--colour", "auto"];   Both => Ok(TerminalColours::Automatic));
    test!(u_never:      TerminalColours <- ["--colour=never"];     Both => Ok(TerminalColours::Never));

    // --color
    test!(no_u_always:  TerminalColours <- ["--color", "always"];  Both => Ok(TerminalColours::Always));
    test!(no_u_auto:    TerminalColours <- ["--color=auto"];       Both => Ok(TerminalColours::Automatic));
    test!(no_u_never:   TerminalColours <- ["--color", "never"];   Both => Ok(TerminalColours::Never));

    // Errors
    test!(no_u_error:   TerminalColours <- ["--color=upstream"];   Both => err Misfire::bad_argument(&flags::COLOR, &os("upstream"), super::COLOURS));  // the error is for --color
    test!(u_error:      TerminalColours <- ["--colour=lovers"];    Both => err Misfire::bad_argument(&flags::COLOR, &os("lovers"),   super::COLOURS));  // and so is this one!

    // Overriding
    test!(overridden_1: TerminalColours <- ["--colour=auto", "--colour=never"];  Last => Ok(TerminalColours::Never));
    test!(overridden_2: TerminalColours <- ["--color=auto",  "--colour=never"];  Last => Ok(TerminalColours::Never));
    test!(overridden_3: TerminalColours <- ["--colour=auto", "--color=never"];   Last => Ok(TerminalColours::Never));
    test!(overridden_4: TerminalColours <- ["--color=auto",  "--color=never"];   Last => Ok(TerminalColours::Never));

    test!(overridden_5: TerminalColours <- ["--colour=auto", "--colour=never"];  Complain => err Misfire::Duplicate(Flag::Long("colour"), Flag::Long("colour")));
    test!(overridden_6: TerminalColours <- ["--color=auto",  "--colour=never"];  Complain => err Misfire::Duplicate(Flag::Long("color"),  Flag::Long("colour")));
    test!(overridden_7: TerminalColours <- ["--colour=auto", "--color=never"];   Complain => err Misfire::Duplicate(Flag::Long("colour"), Flag::Long("color")));
    test!(overridden_8: TerminalColours <- ["--color=auto",  "--color=never"];   Complain => err Misfire::Duplicate(Flag::Long("color"),  Flag::Long("color")));
}

#[cfg(test)]
mod colour_test {
    use super::*;
    use options::flags;
    use options::parser::{Flag, Arg};

    use options::test::parse_for_test;
    use options::test::Strictnesses::*;

    static TEST_ARGS: &[&Arg] = &[ &flags::COLOR,       &flags::COLOUR,
                                   &flags::COLOR_SCALE, &flags::COLOUR_SCALE ];

    macro_rules! test {
        ($name:ident: $type:ident <- $inputs:expr, $widther:expr; $stricts:expr => err $result:expr) => {
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf, &$widther)) {
                    assert_eq!(result.unwrap_err(), $result);
                }
            }
        };

        ($name:ident: $type:ident <- $inputs:expr, $widther:expr; $stricts:expr => like $pat:pat) => {
            #[test]
            fn $name() {
                for result in parse_for_test($inputs.as_ref(), TEST_ARGS, $stricts, |mf| $type::deduce(mf, &$widther)) {
                    println!("Testing {:?}", result);
                    match result {
                        $pat => assert!(true),
                        _    => assert!(false),
                    }
                }
            }
        };
    }

    test!(scale_1: Colours <- ["--color=always", "--color-scale", "--colour-scale"], || None;   Last => like Ok(Colours { scale: true,  .. }));
    test!(scale_2: Colours <- ["--color=always", "--color-scale",                 ], || None;   Last => like Ok(Colours { scale: true,  .. }));
    test!(scale_3: Colours <- ["--color=always",                  "--colour-scale"], || None;   Last => like Ok(Colours { scale: true,  .. }));
    test!(scale_4: Colours <- ["--color=always",                                  ], || None;   Last => like Ok(Colours { scale: false, .. }));

    test!(scale_5: Colours <- ["--color=always", "--color-scale", "--colour-scale"], || None;   Complain => err Misfire::Duplicate(Flag::Long("color-scale"),  Flag::Long("colour-scale")));
    test!(scale_6: Colours <- ["--color=always", "--color-scale",                 ], || None;   Complain => like Ok(Colours { scale: true,  .. }));
    test!(scale_7: Colours <- ["--color=always",                  "--colour-scale"], || None;   Complain => like Ok(Colours { scale: true,  .. }));
    test!(scale_8: Colours <- ["--color=always",                                  ], || None;   Complain => like Ok(Colours { scale: false, .. }));
}
<|endoftext|>


<commit_before>use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::fmt;

use super::parser::FileHeaderParser;
use storage::file_id::*;

use regex::Regex;

#[derive(Debug)]
pub enum FileHeaderSpec {
    Null,
    Bool,
    Integer,
    UInteger,
    Float,
    Text,
    Key { name: &'static str, value_type: Box<FileHeaderSpec> },
    Map { keys: Vec<FileHeaderSpec> },
    Array { allowed_types: Vec<FileHeaderSpec> },
}

#[derive(Debug)]
#[derive(Clone)]
pub enum FileHeaderData {
    Null,
    Bool(bool),
    Integer(i64),
    UInteger(u64),
    Float(f64),
    Text(String),
    Key { name: &'static str, value: Box<FileHeaderData> },
    Map { keys: Vec<FileHeaderData> },
    Array { values: Box<Vec<FileHeaderData>> },
}

impl Display for FileHeaderSpec {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        match self {
            &FileHeaderSpec::Null       => write!(fmt, "NULL"),
            &FileHeaderSpec::Bool       => write!(fmt, "Bool"),
            &FileHeaderSpec::Integer    => write!(fmt, "Integer"),
            &FileHeaderSpec::UInteger   => write!(fmt, "UInteger"),
            &FileHeaderSpec::Float      => write!(fmt, "Float"),
            &FileHeaderSpec::Text       => write!(fmt, "Text"),
            &FileHeaderSpec::Key{name: ref n, value_type: ref vt} => {
                write!(fmt, "Key({:?}) -> {:?}", n, vt)
            }
            &FileHeaderSpec::Map{keys: ref ks} => {
                write!(fmt, "Map -> {:?}", ks)
            }
            &FileHeaderSpec::Array{allowed_types: ref at}  => {
                write!(fmt, "Array({:?})", at)
            }
        }
    }

}

impl FileHeaderData {

    pub fn matches_with(&self, r: &Regex) -> bool {
        match self {
            &FileHeaderData::Text(ref t) => r.is_match(&t[..]),
            &FileHeaderData::Key{name: ref n, value: ref val} => {
                r.is_match(n) || val.matches_with(r)
            },

            &FileHeaderData::Map{keys: ref dks} => {
                dks.iter().any(|x| x.matches_with(r))
            },

            &FileHeaderData::Array{values: ref vs} => {
                vs.iter().any(|x| x.matches_with(r))
            }

            _ => false,
        }
    }
}

pub struct MatchError<'a> {
    summary: String,
    expected: &'a FileHeaderSpec,
    found: &'a FileHeaderData
}

impl<'a> MatchError<'a> {

    pub fn new(s: String,
               ex: &'a FileHeaderSpec,
               found: &'a FileHeaderData) -> MatchError<'a> {
        MatchError {
            summary: s,
            expected: ex,
            found: found,
        }
    }

    pub fn format(&self) -> String {
        format!("MatchError: {:?}\nExpected: {:?}\nFound: {:?}\n",
               self.summary, self.expected, self.found)
    }
}

impl<'a> Error for MatchError<'a> {

    fn description(&self) -> &str {
        &self.summary[..]
    }

    fn cause(&self) -> Option<&Error> {
        None
    }

}

impl<'a> Debug for MatchError<'a> {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "{}", self.format());
        Ok(())
    }

}

impl<'a> Display for MatchError<'a> {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "{}", self.format());
        Ok(())
    }

}

pub fn match_header_spec<'a>(spec: &'a FileHeaderSpec, data: &'a FileHeaderData)
    -> Option<MatchError<'a>>
{
    debug!("Start matching:\n'{:?}'\non\n{:?}", spec, data);
    match (spec, data) {
        (&FileHeaderSpec::Null,     &FileHeaderData::Null)           => { }
        (&FileHeaderSpec::Bool,     &FileHeaderData::Bool(_))        => { }
        (&FileHeaderSpec::Integer,  &FileHeaderData::Integer(_))     => { }
        (&FileHeaderSpec::UInteger, &FileHeaderData::UInteger(_))    => { }
        (&FileHeaderSpec::Float,    &FileHeaderData::Float(_))       => { }
        (&FileHeaderSpec::Text,     &FileHeaderData::Text(_))        => { }

        (
            &FileHeaderSpec::Key{name: ref kname, value_type: ref vtype},
            &FileHeaderData::Key{name: ref n, value: ref val}
        ) => {
            debug!("Matching Key: '{:?}' == '{:?}', Value: '{:?}' == '{:?}'",
                    kname, n,
                    vtype, val);
            if kname != n {
                debug!("Keys not matching");
                unimplemented!();
            }
            return match_header_spec(&*vtype, &*val);
        }

        (
            &FileHeaderSpec::Map{keys: ref sks},
            &FileHeaderData::Map{keys: ref dks}
        ) => {
            debug!("Matching Map: '{:?}' == '{:?}'", sks, dks);

            for (s, d) in sks.iter().zip(dks.iter()) {
                let res = match_header_spec(s, d);
                if res.is_some() {
                    return res;
                }
            }
        }

        (
            &FileHeaderSpec::Array{allowed_types: ref vtypes},
            &FileHeaderData::Array{values: ref vs}
        ) => {
            debug!("Matching Array: '{:?}' == '{:?}'", vtypes, vs);
            for (t, v) in vtypes.iter().zip(vs.iter()) {
                let res = match_header_spec(t, v);
                if res.is_some() {
                    return res;
                }
            }
        }

        (k, v) => {
            return Some(MatchError::new(String::from("Expected type does not match found type"),
                                 k, v
                                 ))
        }
    }
    None
}

/*
 * Internal abstract view on a file. Does not exist on the FS and is just kept
 * internally until it is written to disk.
 */
pub struct File {
    header  : FileHeaderData,
    data    : String,
    id      : FileID,
}

impl File {

    pub fn new() -> File {
        let f = File {
            header: FileHeaderData::Null,
            data: String::from(""),
            id: File::get_new_file_id(),
        };
        debug!("Create new File object: {:?}", f);
        f
    }

    pub fn from_parser_result(id: FileID, header: FileHeaderData, data: String) -> File {
        let f = File {
            header: header,
            data: data,
            id: id,
        };
        debug!("Create new File object from parser result: {:?}", f);
        f
    }

    pub fn new_with_header(h: FileHeaderData) -> File {
        let f = File {
            header: h,
            data: String::from(""),
            id: File::get_new_file_id(),
        };
        debug!("Create new File object with header: {:?}", f);
        f
    }

    pub fn new_with_data(d: String) -> File {
        let f = File {
            header: FileHeaderData::Null,
            data: d,
            id: File::get_new_file_id(),
        };
        debug!("Create new File object with data: {:?}", f);
        f
    }

    pub fn new_with_content(h: FileHeaderData, d: String) -> File {
        let f = File {
            header: h,
            data: d,
            id: File::get_new_file_id(),
        };
        debug!("Create new File object with content: {:?}", f);
        f
    }

    pub fn contents(&self) -> (FileHeaderData, String) {
        (self.header.clone(), self.data.clone())
    }

    pub fn id(&self) -> FileID {
        self.id.clone()
    }

    pub fn matches_with(&self, r: &Regex) -> bool {
        r.is_match(&self.data[..]) || self.header.matches_with(r)
    }

    fn get_new_file_id() -> FileID {
        use uuid::Uuid;
        Uuid::new_v4().to_hyphenated_string()
    }
}

impl Debug for File {

    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "File[{:?}] header: '{:?}', data: '{:?}')",
            self.id, self.header, self.data)
    }
}

<commit_msg>Implement File::{header(),data()} as header/data getters<commit_after>use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
use std::fmt;

use super::parser::FileHeaderParser;
use storage::file_id::*;

use regex::Regex;

#[derive(Debug)]
pub enum FileHeaderSpec {
    Null,
    Bool,
    Integer,
    UInteger,
    Float,
    Text,
    Key { name: &'static str, value_type: Box<FileHeaderSpec> },
    Map { keys: Vec<FileHeaderSpec> },
    Array { allowed_types: Vec<FileHeaderSpec> },
}

#[derive(Debug)]
#[derive(Clone)]
pub enum FileHeaderData {
    Null,
    Bool(bool),
    Integer(i64),
    UInteger(u64),
    Float(f64),
    Text(String),
    Key { name: &'static str, value: Box<FileHeaderData> },
    Map { keys: Vec<FileHeaderData> },
    Array { values: Box<Vec<FileHeaderData>> },
}

impl Display for FileHeaderSpec {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        match self {
            &FileHeaderSpec::Null       => write!(fmt, "NULL"),
            &FileHeaderSpec::Bool       => write!(fmt, "Bool"),
            &FileHeaderSpec::Integer    => write!(fmt, "Integer"),
            &FileHeaderSpec::UInteger   => write!(fmt, "UInteger"),
            &FileHeaderSpec::Float      => write!(fmt, "Float"),
            &FileHeaderSpec::Text       => write!(fmt, "Text"),
            &FileHeaderSpec::Key{name: ref n, value_type: ref vt} => {
                write!(fmt, "Key({:?}) -> {:?}", n, vt)
            }
            &FileHeaderSpec::Map{keys: ref ks} => {
                write!(fmt, "Map -> {:?}", ks)
            }
            &FileHeaderSpec::Array{allowed_types: ref at}  => {
                write!(fmt, "Array({:?})", at)
            }
        }
    }

}

impl FileHeaderData {

    pub fn matches_with(&self, r: &Regex) -> bool {
        match self {
            &FileHeaderData::Text(ref t) => r.is_match(&t[..]),
            &FileHeaderData::Key{name: ref n, value: ref val} => {
                r.is_match(n) || val.matches_with(r)
            },

            &FileHeaderData::Map{keys: ref dks} => {
                dks.iter().any(|x| x.matches_with(r))
            },

            &FileHeaderData::Array{values: ref vs} => {
                vs.iter().any(|x| x.matches_with(r))
            }

            _ => false,
        }
    }
}

pub struct MatchError<'a> {
    summary: String,
    expected: &'a FileHeaderSpec,
    found: &'a FileHeaderData
}

impl<'a> MatchError<'a> {

    pub fn new(s: String,
               ex: &'a FileHeaderSpec,
               found: &'a FileHeaderData) -> MatchError<'a> {
        MatchError {
            summary: s,
            expected: ex,
            found: found,
        }
    }

    pub fn format(&self) -> String {
        format!("MatchError: {:?}\nExpected: {:?}\nFound: {:?}\n",
               self.summary, self.expected, self.found)
    }
}

impl<'a> Error for MatchError<'a> {

    fn description(&self) -> &str {
        &self.summary[..]
    }

    fn cause(&self) -> Option<&Error> {
        None
    }

}

impl<'a> Debug for MatchError<'a> {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "{}", self.format());
        Ok(())
    }

}

impl<'a> Display for MatchError<'a> {

    fn fmt(&self, fmt: &mut Formatter) -> fmt::Result {
        write!(fmt, "{}", self.format());
        Ok(())
    }

}

pub fn match_header_spec<'a>(spec: &'a FileHeaderSpec, data: &'a FileHeaderData)
    -> Option<MatchError<'a>>
{
    debug!("Start matching:\n'{:?}'\non\n{:?}", spec, data);
    match (spec, data) {
        (&FileHeaderSpec::Null,     &FileHeaderData::Null)           => { }
        (&FileHeaderSpec::Bool,     &FileHeaderData::Bool(_))        => { }
        (&FileHeaderSpec::Integer,  &FileHeaderData::Integer(_))     => { }
        (&FileHeaderSpec::UInteger, &FileHeaderData::UInteger(_))    => { }
        (&FileHeaderSpec::Float,    &FileHeaderData::Float(_))       => { }
        (&FileHeaderSpec::Text,     &FileHeaderData::Text(_))        => { }

        (
            &FileHeaderSpec::Key{name: ref kname, value_type: ref vtype},
            &FileHeaderData::Key{name: ref n, value: ref val}
        ) => {
            debug!("Matching Key: '{:?}' == '{:?}', Value: '{:?}' == '{:?}'",
                    kname, n,
                    vtype, val);
            if kname != n {
                debug!("Keys not matching");
                unimplemented!();
            }
            return match_header_spec(&*vtype, &*val);
        }

        (
            &FileHeaderSpec::Map{keys: ref sks},
            &FileHeaderData::Map{keys: ref dks}
        ) => {
            debug!("Matching Map: '{:?}' == '{:?}'", sks, dks);

            for (s, d) in sks.iter().zip(dks.iter()) {
                let res = match_header_spec(s, d);
                if res.is_some() {
                    return res;
                }
            }
        }

        (
            &FileHeaderSpec::Array{allowed_types: ref vtypes},
            &FileHeaderData::Array{values: ref vs}
        ) => {
            debug!("Matching Array: '{:?}' == '{:?}'", vtypes, vs);
            for (t, v) in vtypes.iter().zip(vs.iter()) {
                let res = match_header_spec(t, v);
                if res.is_some() {
                    return res;
                }
            }
        }

        (k, v) => {
            return Some(MatchError::new(String::from("Expected type does not match found type"),
                                 k, v
                                 ))
        }
    }
    None
}

/*
 * Internal abstract view on a file. Does not exist on the FS and is just kept
 * internally until it is written to disk.
 */
pub struct File {
    header  : FileHeaderData,
    data    : String,
    id      : FileID,
}

impl File {

    pub fn new() -> File {
        let f = File {
            header: FileHeaderData::Null,
            data: String::from(""),
            id: File::get_new_file_id(),
        };
        debug!("Create new File object: {:?}", f);
        f
    }

    pub fn from_parser_result(id: FileID, header: FileHeaderData, data: String) -> File {
        let f = File {
            header: header,
            data: data,
            id: id,
        };
        debug!("Create new File object from parser result: {:?}", f);
        f
    }

    pub fn new_with_header(h: FileHeaderData) -> File {
        let f = File {
            header: h,
            data: String::from(""),
            id: File::get_new_file_id(),
        };
        debug!("Create new File object with header: {:?}", f);
        f
    }

    pub fn new_with_data(d: String) -> File {
        let f = File {
            header: FileHeaderData::Null,
            data: d,
            id: File::get_new_file_id(),
        };
        debug!("Create new File object with data: {:?}", f);
        f
    }

    pub fn new_with_content(h: FileHeaderData, d: String) -> File {
        let f = File {
            header: h,
            data: d,
            id: File::get_new_file_id(),
        };
        debug!("Create new File object with content: {:?}", f);
        f
    }

    pub fn header(&self) -> FileHeaderData {
        self.header.clone()
    }

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

    pub fn contents(&self) -> (FileHeaderData, String) {
        (self.header.clone(), self.data.clone())
    }

    pub fn id(&self) -> FileID {
        self.id.clone()
    }

    pub fn matches_with(&self, r: &Regex) -> bool {
        r.is_match(&self.data[..]) || self.header.matches_with(r)
    }

    fn get_new_file_id() -> FileID {
        use uuid::Uuid;
        Uuid::new_v4().to_hyphenated_string()
    }
}

impl Debug for File {

    fn fmt(&self, f: &mut fmt::Formatter) -> Result<(), fmt::Error> {
        write!(f, "File[{:?}] header: '{:?}', data: '{:?}')",
            self.id, self.header, self.data)
    }
}

<|endoftext|>


<commit_before>use std::env;
use std::error;
use std::ffi::{OsStr, OsString};
use std::fmt;
use std::fs::File;
use std::io;
use std::io::Read;
use std::path::PathBuf;
use std::slice;
use std::str;

pub mod cap;

pub struct Desc {
    pub names: Vec<String>,
    pub bools: Vec<bool>,
    pub nums: Vec<u16>,
    pub strings: Vec<String>,
}

impl Desc {
    pub fn file<S: AsRef<OsStr>>(name: S) -> Option<File> {
        let name = name.as_ref();

        let first_char = match name.to_str() {
            Some(s) if !s.is_empty() => s.chars().next().unwrap(),
            _ => return None,
        };
        let first_hex = format!("{:x}", first_char as usize);
        let first_char = first_char.to_string();

        let d1 = to_path(env::var_os("TERMINFO"));
        let d2 = to_path(env::home_dir()).map(|d| d.join(".terminfo"));
        let d3 = to_paths(env::var_os("TERMINFO_DIRS"));
        let d3 = d3.into_iter().map(|d| if d.as_os_str().is_empty() {
                                        PathBuf::from("/usr/share/terminfo")
                                    } else {
                                        d
                                    });
        let d4 = vec![PathBuf::from("/etc/terminfo"),
                      PathBuf::from("/lib/terminfo"),
                      PathBuf::from("/usr/share/terminfo")];
        let ds = d1.into_iter()
            .chain(d2.into_iter())
            .chain(d3)
            .chain(d4.into_iter());

        for d in ds {
            if let Ok(f) = File::open(d.join(&first_char).join(name)) {
                return Some(f);
            }
            if let Ok(f) = File::open(d.join(&first_hex).join(name)) {
                return Some(f);
            }
        }
        return None;
    }

    pub fn parse(r: &mut Read) -> Result<Desc, ParseError> {
        let header = read_words(r, 6)?;
        if header[0] != 282 {
            return Err(parse_error("wrong magic number"));
        }

        let name_sz = header[1] as usize;
        if name_sz == 0 {
            return Err(parse_error("zero-length name section"));
        }
        let name_buf = read_bytes(r, name_sz)?;
        if name_buf[name_sz - 1] != 0 {
            return Err(parse_error("name section is not null-terminated"));
        }
        let names = str::from_utf8(&name_buf[0..name_sz - 1])?;
        let names: Vec<_> = names.split('|').map(str::to_owned).collect();
        if names.len() < 2 {
            return Err(parse_error("too few items in name section"));
        }

        let bools_num = header[2] as usize;
        if bools_num > cap::NUM_BOOLS {
            return Err(parse_error("too many boolean flags"));
        }
        let bool_buf = read_bytes(r, bools_num)?;
        let bools = bool_buf.into_iter()
            .map(|b| if b == 0 { false } else { true })
            .collect();

        if (name_sz + bools_num) % 2 != 0 {
            read_bytes(r, 1)?;
        }

        let ints_num = header[3] as usize;
        if ints_num > cap::NUM_INTS {
            return Err(parse_error("too many numbers"));
        }
        let nums = read_words(r, ints_num)?;

        let strings_num = header[4] as usize;
        let string_sz = header[5] as usize;
        if strings_num > cap::NUM_STRINGS {
            return Err(parse_error("too many strings"));
        }
        let offsets = read_words(r, strings_num)?;
        let str_table = read_bytes(r, string_sz)?;
        let mut strings = Vec::new();
        for pos in offsets {
            let pos = pos as usize;
            if pos == 0xffff {
                strings.push(String::new());
            } else if pos >= string_sz {
                return Err(parse_error("invalid string offset"));
            } else {
                match str_table[pos..].iter().position(|&b| b == 0) {
                    None => return Err(parse_error("unterminated string")),
                    Some(end) => {
                        let s = str::from_utf8(&str_table[pos..pos + end])?;
                        strings.push(s.to_owned());
                    }
                }
            }
        }

        Ok(Desc {
               names: names,
               bools: bools,
               nums: nums,
               strings: strings,
           })
    }
}

#[derive(Debug)]
pub struct ParseError {
    inner: ParseErrorImpl,
}

fn parse_error(msg: &str) -> ParseError {
    ParseError { inner: ParseErrorImpl::Other(msg.to_owned()) }
}

#[derive(Debug)]
enum ParseErrorImpl {
    Io(io::Error),
    Other(String),
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.inner {
            ParseErrorImpl::Io(ref err) => err.fmt(f),
            ParseErrorImpl::Other(ref msg) => write!(f, "{}", msg),
        }
    }
}

impl error::Error for ParseError {
    fn description(&self) -> &str {
        match self.inner {
            ParseErrorImpl::Io(ref err) => err.description(),
            _ => "invalid terminfo description",
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match self.inner {
            ParseErrorImpl::Io(ref err) => err.cause(),
            _ => None,
        }
    }
}

impl From<io::Error> for ParseError {
    fn from(err: io::Error) -> ParseError {
        ParseError { inner: ParseErrorImpl::Io(err) }
    }
}

impl From<str::Utf8Error> for ParseError {
    fn from(err: str::Utf8Error) -> ParseError {
        parse_error(&err.to_string())
    }
}

fn to_path<T: Into<PathBuf>>(var: Option<T>) -> Option<PathBuf> {
    match var {
        Some(d) => {
            let d = d.into();
            if d.as_os_str().is_empty() {
                None
            } else {
                Some(d)
            }
        }
        None => None,
    }
}

fn to_paths(var: Option<OsString>) -> Vec<PathBuf> {
    match var {
        None => Vec::new(),
        Some(ref d) if d.is_empty() => Vec::new(),
        Some(d) => env::split_paths(&d).collect(),
    }
}

fn read_bytes(r: &mut Read, n: usize) -> io::Result<Vec<u8>> {
    let mut buf = Vec::with_capacity(n);
    if r.take(n as u64).read_to_end(&mut buf)? < n {
        return Err(io::Error::new(io::ErrorKind::UnexpectedEof,
                                  "failed to fill whole buffer"));
    }
    Ok(buf)
}

fn read_words(r: &mut Read, n: usize) -> io::Result<Vec<u16>> {
    let mut buf_16: Vec<u16> = vec![0; n];
    let mut buf_8 = unsafe {
        slice::from_raw_parts_mut(buf_16.as_mut_ptr() as *mut u8, n * 2)
    };
    r.read_exact(&mut buf_8)?;
    if cfg!(target_endian = "big") {
        for w in &mut buf_16 {
            *w = w.swap_bytes();
        }
    }
    Ok(buf_16)
}
<commit_msg>Desc::file takes &str instead of AsRef<OsString><commit_after>use std::env;
use std::error;
use std::fmt;
use std::fs::File;
use std::io;
use std::io::Read;
use std::path::PathBuf;
use std::slice;
use std::str;

pub mod cap;

pub struct Desc {
    pub names: Vec<String>,
    pub bools: Vec<bool>,
    pub nums: Vec<u16>,
    pub strings: Vec<String>,
}

impl Desc {
    pub fn file(name: &str) -> Option<File> {
        if name.is_empty() {
            return None;
        }

        let first_char = name.chars().next().unwrap();
        let first_hex = format!("{:x}", first_char as usize);
        let first_char = first_char.to_string();

        let d1 = to_path(env::var("TERMINFO").ok());
        let d2 = to_path(env::home_dir()).map(|d| d.join(".terminfo"));
        let d3 = to_paths(env::var("TERMINFO_DIRS").ok());
        let d3 = d3.into_iter().map(|d| if d.as_os_str().is_empty() {
                                        PathBuf::from("/usr/share/terminfo")
                                    } else {
                                        d
                                    });
        let d4 = vec![PathBuf::from("/etc/terminfo"),
                      PathBuf::from("/lib/terminfo"),
                      PathBuf::from("/usr/share/terminfo")];
        let ds = d1.into_iter()
            .chain(d2.into_iter())
            .chain(d3)
            .chain(d4.into_iter());

        for d in ds {
            if let Ok(f) = File::open(d.join(&first_char).join(name)) {
                return Some(f);
            }
            if let Ok(f) = File::open(d.join(&first_hex).join(name)) {
                return Some(f);
            }
        }
        return None;
    }

    pub fn parse(r: &mut Read) -> Result<Desc, ParseError> {
        let header = read_words(r, 6)?;
        if header[0] != 282 {
            return Err(parse_error("wrong magic number"));
        }

        let name_sz = header[1] as usize;
        if name_sz == 0 {
            return Err(parse_error("zero-length name section"));
        }
        let name_buf = read_bytes(r, name_sz)?;
        if name_buf[name_sz - 1] != 0 {
            return Err(parse_error("name section is not null-terminated"));
        }
        let names = str::from_utf8(&name_buf[0..name_sz - 1])?;
        let names: Vec<_> = names.split('|').map(str::to_owned).collect();
        if names.len() < 2 {
            return Err(parse_error("too few items in name section"));
        }

        let bools_num = header[2] as usize;
        if bools_num > cap::NUM_BOOLS {
            return Err(parse_error("too many boolean flags"));
        }
        let bool_buf = read_bytes(r, bools_num)?;
        let bools = bool_buf.into_iter()
            .map(|b| if b == 0 { false } else { true })
            .collect();

        if (name_sz + bools_num) % 2 != 0 {
            read_bytes(r, 1)?;
        }

        let ints_num = header[3] as usize;
        if ints_num > cap::NUM_INTS {
            return Err(parse_error("too many numbers"));
        }
        let nums = read_words(r, ints_num)?;

        let strings_num = header[4] as usize;
        let string_sz = header[5] as usize;
        if strings_num > cap::NUM_STRINGS {
            return Err(parse_error("too many strings"));
        }
        let offsets = read_words(r, strings_num)?;
        let str_table = read_bytes(r, string_sz)?;
        let mut strings = Vec::new();
        for pos in offsets {
            let pos = pos as usize;
            if pos == 0xffff {
                strings.push(String::new());
            } else if pos >= string_sz {
                return Err(parse_error("invalid string offset"));
            } else {
                match str_table[pos..].iter().position(|&b| b == 0) {
                    None => return Err(parse_error("unterminated string")),
                    Some(end) => {
                        let s = str::from_utf8(&str_table[pos..pos + end])?;
                        strings.push(s.to_owned());
                    }
                }
            }
        }

        Ok(Desc {
               names: names,
               bools: bools,
               nums: nums,
               strings: strings,
           })
    }
}

#[derive(Debug)]
pub struct ParseError {
    inner: ParseErrorImpl,
}

fn parse_error(msg: &str) -> ParseError {
    ParseError { inner: ParseErrorImpl::Other(msg.to_owned()) }
}

#[derive(Debug)]
enum ParseErrorImpl {
    Io(io::Error),
    Other(String),
}

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self.inner {
            ParseErrorImpl::Io(ref err) => err.fmt(f),
            ParseErrorImpl::Other(ref msg) => write!(f, "{}", msg),
        }
    }
}

impl error::Error for ParseError {
    fn description(&self) -> &str {
        match self.inner {
            ParseErrorImpl::Io(ref err) => err.description(),
            _ => "invalid terminfo description",
        }
    }

    fn cause(&self) -> Option<&error::Error> {
        match self.inner {
            ParseErrorImpl::Io(ref err) => err.cause(),
            _ => None,
        }
    }
}

impl From<io::Error> for ParseError {
    fn from(err: io::Error) -> ParseError {
        ParseError { inner: ParseErrorImpl::Io(err) }
    }
}

impl From<str::Utf8Error> for ParseError {
    fn from(err: str::Utf8Error) -> ParseError {
        parse_error(&err.to_string())
    }
}

fn to_path<T: Into<PathBuf>>(var: Option<T>) -> Option<PathBuf> {
    match var {
        Some(d) => {
            let d = d.into();
            if d.as_os_str().is_empty() {
                None
            } else {
                Some(d)
            }
        }
        None => None,
    }
}

fn to_paths(var: Option<String>) -> Vec<PathBuf> {
    match var {
        None => Vec::new(),
        Some(ref d) if d.is_empty() => Vec::new(),
        Some(d) => env::split_paths(&d).collect(),
    }
}

fn read_bytes(r: &mut Read, n: usize) -> io::Result<Vec<u8>> {
    let mut buf = Vec::with_capacity(n);
    if r.take(n as u64).read_to_end(&mut buf)? < n {
        return Err(io::Error::new(io::ErrorKind::UnexpectedEof,
                                  "failed to fill whole buffer"));
    }
    Ok(buf)
}

fn read_words(r: &mut Read, n: usize) -> io::Result<Vec<u16>> {
    let mut buf_16: Vec<u16> = vec![0; n];
    let mut buf_8 = unsafe {
        slice::from_raw_parts_mut(buf_16.as_mut_ptr() as *mut u8, n * 2)
    };
    r.read_exact(&mut buf_8)?;
    if cfg!(target_endian = "big") {
        for w in &mut buf_16 {
            *w = w.swap_bytes();
        }
    }
    Ok(buf_16)
}
<|endoftext|>