<commit_before>//! choice combinators

#[macro_use]
mod macros;

use crate::error::ErrorKind;
use crate::error::ParseError;
use crate::internal::{Err, IResult};

/// helper trait for the [alt()] combinator
///
/// this trait is implemented for tuples of up to 21 elements
pub trait Alt<I, O, E> {
  /// tests each parser in the tuple and returns the result of the first one that succeeds
  fn choice(&self, input: I) -> IResult<I, O, E>;
}

/// tests a list of parsers one by one until one succeeds
///
/// It takes as argument a tuple of parsers.
///
/// ```rust
/// # #[macro_use] extern crate nom;
/// # use nom::{Err,error::ErrorKind, Needed, IResult};
/// use nom::character::complete::{alpha1, digit1};
/// use nom::branch::alt;
/// # fn main() {
/// fn parser(input: &str) -> IResult<&str, &str> {
///   alt((alpha1, digit1))(input)
/// };
///
/// // the first parser, alpha1, recognizes the input
/// assert_eq!(parser("abc"), Ok(("", "abc")));
///
/// // the first parser returns an error, so alt tries the second one
/// assert_eq!(parser("123456"), Ok(("", "123456")));
///
/// // both parsers failed, and with the default error type, alt will return the last error
/// assert_eq!(parser(" "), Err(Err::Error(error_position!(" ", ErrorKind::Digit))));
/// # }
/// ```
///
/// with a custom error type, it is possible to have alt return the error of the parser
/// that went the farthest in the input data
pub fn alt<I: Clone, O, E: ParseError<I>, List: Alt<I, O, E>>(l: List) -> impl Fn(I) -> IResult<I, O, E> {
  move |i: I| l.choice(i)
}

/// helper trait for the [permutation()] combinator
///
/// this trait is implemented for tuples of up to 21 elements
pub trait Permutation<I, O, E> {
  /// tries to apply all parsers in the tuple in various orders until all of them succeed
  fn permutation(&self, input: I) -> IResult<I, O, E>;
}

/// applies a list of parsers in any order
///
/// permutation will succeed if all of the child parsers succeeded.
/// It takes as argument a tuple of parsers, and returns a
/// tuple of the parser results.
///
/// ```rust
/// # #[macro_use] extern crate nom;
/// # use nom::{Err,error::ErrorKind, Needed, IResult};
/// use nom::character::complete::{alpha1, digit1};
/// use nom::branch::permutation;
/// # fn main() {
/// fn parser(input: &str) -> IResult<&str, (&str, &str)> {
///   permutation((alpha1, digit1))(input)
/// };
///
/// // permutation recognizes alphabetic characters then digit
/// assert_eq!(parser("abc123"), Ok(("", ("abc", "123"))));
///
/// // but also in inverse order
/// assert_eq!(parser("123abc"), Ok(("", ("abc", "123"))));
///
/// // it will fail if one of the parsers failed
/// assert_eq!(parser("abc;"), Err(Err::Error(error_position!(";", ErrorKind::Permutation))));
/// # }
/// ```
pub fn permutation<I: Clone, O, E: ParseError<I>, List: Permutation<I, O, E>>(l: List) -> impl Fn(I) -> IResult<I, O, E> {
  move |i: I| l.permutation(i)
}

macro_rules! alt_trait(
  ($first:ident $second:ident $($id: ident)+) => (
    alt_trait!(__impl $first $second; $($id)+);
  );
  (__impl $($current:ident)*; $head:ident $($id: ident)+) => (
    alt_trait_impl!($($current)*);

    alt_trait!(__impl $($current)* $head; $($id)+);
  );
  (__impl $($current:ident)*; $head:ident) => (
    alt_trait_impl!($($current)*);
    alt_trait_impl!($($current)* $head);
  );
);

macro_rules! alt_trait_impl(
  ($($id:ident)+) => (
    impl<
      Input: Clone, Output, Error: ParseError<Input>,
      $($id: Fn(Input) -> IResult<Input, Output, Error>),+
    > Alt<Input, Output, Error> for ( $($id),+ ) {

      fn choice(&self, input: Input) -> IResult<Input, Output, Error> {
        let mut err = None;
        alt_trait_inner!(0, self, input, err, $($id)+);

        Err(Err::Error(Error::append(input, ErrorKind::Alt, err.unwrap())))
      }
    }
  );
);

macro_rules! alt_trait_inner(
  ($it:tt, $self:expr, $input:expr, $err:expr, $head:ident $($id:ident)+) => (
    match $self.$it($input.clone()) {
      Err(Err::Error(e)) => {
        if $err.is_none() {
          $err = Some(e);
        } else {
          $err = Some($err.unwrap().or(e));
        }
        succ!($it, alt_trait_inner!($self, $input, $err, $($id)+))
      },
      res => return res,
    }
  );
  ($it:tt, $self:expr, $input:expr, $err:expr, $head:ident) => (
    match $self.$it($input.clone()) {
      Err(Err::Error(e)) => {
        if $err.is_none() {
          $err = Some(e);
        } else {
          $err = Some($err.unwrap().or(e));
        }
      },
      res => return res,
    }
  );
);

alt_trait!(A B C D E F G H I J K L M N O P Q R S T U);

macro_rules! permutation_trait(
  ($name1:ident $ty1:ident, $name2:ident $ty2:ident) => (
    permutation_trait_impl!($name1 $ty1, $name2 $ty2);
  );
  ($name1:ident $ty1:ident, $name2: ident $ty2:ident, $($name:ident $ty:ident),*) => (
    permutation_trait!(__impl $name1 $ty1, $name2 $ty2; $($name $ty),*);
  );
  (__impl $($name:ident $ty: ident),+; $name1:ident $ty1:ident, $($name2:ident $ty2:ident),*) => (
    permutation_trait_impl!($($name $ty),+);
    permutation_trait!(__impl $($name $ty),+ , $name1 $ty1; $($name2 $ty2),*);
  );
  (__impl $($name:ident $ty: ident),+; $name1:ident $ty1:ident) => (
    permutation_trait_impl!($($name $ty),+);
    permutation_trait_impl!($($name $ty),+, $name1 $ty1);
  );
);

macro_rules! permutation_trait_impl(
  ($($name:ident $ty: ident),+) => (
    impl<
      Input: Clone, $($ty),+ , Error: ParseError<Input>,
      $($name: Fn(Input) -> IResult<Input, $ty, Error>),+
    > Permutation<Input, ( $($ty),+ ), Error> for ( $($name),+ ) {

      fn permutation(&self, mut input: Input) -> IResult<Input, ( $($ty),+ ), Error> {
        let mut res = permutation_init!((), $($name),+);

        loop {
          let mut all_done = true;
          permutation_trait_inner!(0, self, input, res, all_done, $($name)+);

          //if we reach that part, it means none of the parsers were able to read anything
          if !all_done {
            //FIXME: should wrap the error returned by the child parser
            return Err(Err::Error(error_position!(input, ErrorKind::Permutation)));
          }
          break;
        }

        if let Some(unwrapped_res) = { permutation_trait_unwrap!(0, (), res, $($name),+) } {
          Ok((input, unwrapped_res))
        } else {
          Err(Err::Error(error_position!(input, ErrorKind::Permutation)))
        }
      }
    }
  );
);

macro_rules! permutation_trait_inner(
  ($it:tt, $self:expr, $input:ident, $res:expr, $all_done:expr, $head:ident $($id:ident)+) => ({
    if $res.$it.is_none() {
      match $self.$it($input.clone()) {
        Ok((i,o))     => {
          $input = i;
          $res.$it = Some(o);
          continue;
        },
        Err(Err::Error(_)) => {
          $all_done = false;
        },
        Err(e) => {
          return Err(e);
        }
      };
    }
    succ!($it, permutation_trait_inner!($self, $input, $res, $all_done, $($id)+));
  });
  ($it:tt, $self:expr, $input:ident, $res:expr, $all_done:expr, $head:ident) => ({
    if $res.$it.is_none() {
      match $self.$it($input.clone()) {
        Ok((i,o))     => {
          $input = i;
          $res.$it = Some(o);
          continue;
        },
        Err(Err::Error(_)) => {
          $all_done = false;
        },
        Err(e) => {
          return Err(e);
        }
      };
    }
  });
);

macro_rules! permutation_trait_unwrap (
  ($it:tt,  (), $res:ident, $e:ident, $($name:ident),+) => ({
    let res = $res.$it;
    if res.is_some() {
      succ!($it, permutation_trait_unwrap!((res.unwrap()), $res, $($name),+))
    } else {
      $crate::lib::std::option::Option::None
    }
  });

  ($it:tt, ($($parsed:expr),*), $res:ident, $e:ident, $($name:ident),+) => ({
    let res = $res.$it;
    if res.is_some() {
      succ!($it, permutation_trait_unwrap!(($($parsed),* , res.unwrap()), $res, $($name),+))
    } else {
      $crate::lib::std::option::Option::None
    }
  });

  ($it:tt, ($($parsed:expr),*), $res:ident, $name:ident) => ({
    let res = $res.$it;
    if res.is_some() {
      $crate::lib::std::option::Option::Some(($($parsed),* , res.unwrap() ))
    } else {
      $crate::lib::std::option::Option::None
    }
  });
);

permutation_trait!(FnA A, FnB B, FnC C, FnD D, FnE E, FnF F, FnG G, FnH H, FnI I, FnJ J, FnK K, FnL L, FnM M, FnN N, FnO O, FnP P, FnQ Q, FnR R, FnS S, FnT T, FnU U);
<commit_msg>Simplified alt() implementation<commit_after>//! choice combinators

#[macro_use]
mod macros;

use crate::error::ErrorKind;
use crate::error::ParseError;
use crate::internal::{Err, IResult};

/// helper trait for the [alt()] combinator
///
/// this trait is implemented for tuples of up to 21 elements
pub trait Alt<I, O, E> {
  /// tests each parser in the tuple and returns the result of the first one that succeeds
  fn choice(&self, input: I) -> IResult<I, O, E>;
}

/// tests a list of parsers one by one until one succeeds
///
/// It takes as argument a tuple of parsers.
///
/// ```rust
/// # #[macro_use] extern crate nom;
/// # use nom::{Err,error::ErrorKind, Needed, IResult};
/// use nom::character::complete::{alpha1, digit1};
/// use nom::branch::alt;
/// # fn main() {
/// fn parser(input: &str) -> IResult<&str, &str> {
///   alt((alpha1, digit1))(input)
/// };
///
/// // the first parser, alpha1, recognizes the input
/// assert_eq!(parser("abc"), Ok(("", "abc")));
///
/// // the first parser returns an error, so alt tries the second one
/// assert_eq!(parser("123456"), Ok(("", "123456")));
///
/// // both parsers failed, and with the default error type, alt will return the last error
/// assert_eq!(parser(" "), Err(Err::Error(error_position!(" ", ErrorKind::Digit))));
/// # }
/// ```
///
/// with a custom error type, it is possible to have alt return the error of the parser
/// that went the farthest in the input data
pub fn alt<I: Clone, O, E: ParseError<I>, List: Alt<I, O, E>>(l: List) -> impl Fn(I) -> IResult<I, O, E> {
  move |i: I| l.choice(i)
}

/// helper trait for the [permutation()] combinator
///
/// this trait is implemented for tuples of up to 21 elements
pub trait Permutation<I, O, E> {
  /// tries to apply all parsers in the tuple in various orders until all of them succeed
  fn permutation(&self, input: I) -> IResult<I, O, E>;
}

/// applies a list of parsers in any order
///
/// permutation will succeed if all of the child parsers succeeded.
/// It takes as argument a tuple of parsers, and returns a
/// tuple of the parser results.
///
/// ```rust
/// # #[macro_use] extern crate nom;
/// # use nom::{Err,error::ErrorKind, Needed, IResult};
/// use nom::character::complete::{alpha1, digit1};
/// use nom::branch::permutation;
/// # fn main() {
/// fn parser(input: &str) -> IResult<&str, (&str, &str)> {
///   permutation((alpha1, digit1))(input)
/// };
///
/// // permutation recognizes alphabetic characters then digit
/// assert_eq!(parser("abc123"), Ok(("", ("abc", "123"))));
///
/// // but also in inverse order
/// assert_eq!(parser("123abc"), Ok(("", ("abc", "123"))));
///
/// // it will fail if one of the parsers failed
/// assert_eq!(parser("abc;"), Err(Err::Error(error_position!(";", ErrorKind::Permutation))));
/// # }
/// ```
pub fn permutation<I: Clone, O, E: ParseError<I>, List: Permutation<I, O, E>>(l: List) -> impl Fn(I) -> IResult<I, O, E> {
  move |i: I| l.permutation(i)
}

macro_rules! alt_trait(
  ($first:ident $second:ident $($id: ident)+) => (
    alt_trait!(__impl $first $second; $($id)+);
  );
  (__impl $($current:ident)*; $head:ident $($id: ident)+) => (
    alt_trait_impl!($($current)*);

    alt_trait!(__impl $($current)* $head; $($id)+);
  );
  (__impl $($current:ident)*; $head:ident) => (
    alt_trait_impl!($($current)*);
    alt_trait_impl!($($current)* $head);
  );
);

macro_rules! alt_trait_impl(
  ($($id:ident)+) => (
    impl<
      Input: Clone, Output, Error: ParseError<Input>,
      $($id: Fn(Input) -> IResult<Input, Output, Error>),+
    > Alt<Input, Output, Error> for ( $($id),+ ) {

      fn choice(&self, input: Input) -> IResult<Input, Output, Error> {
        let mut err: Option<Error> = None;
        alt_trait_inner!(0, self, input, err, $($id)+);

        Err(Err::Error(Error::append(input, ErrorKind::Alt, err.unwrap())))
      }
    }
  );
);

macro_rules! alt_trait_inner(
  ($it:tt, $self:expr, $input:expr, $err:expr, $head:ident $($id:ident)+) => (
    match $self.$it($input.clone()) {
      Err(Err::Error(e)) => {
        $err = Some(match $err.take() {
          None => e,
          Some(prev) => prev.or(e),
        });
        succ!($it, alt_trait_inner!($self, $input, $err, $($id)+))
      },
      res => return res,
    }
  );
  ($it:tt, $self:expr, $input:expr, $err:expr, $head:ident) => (
    match $self.$it($input.clone()) {
      Err(Err::Error(e)) => {
        $err = Some(match $err.take() {
          None => e,
          Some(prev) => prev.or(e),
        });
      },
      res => return res,
    }
  );
);

alt_trait!(A B C D E F G H I J K L M N O P Q R S T U);

macro_rules! permutation_trait(
  ($name1:ident $ty1:ident, $name2:ident $ty2:ident) => (
    permutation_trait_impl!($name1 $ty1, $name2 $ty2);
  );
  ($name1:ident $ty1:ident, $name2: ident $ty2:ident, $($name:ident $ty:ident),*) => (
    permutation_trait!(__impl $name1 $ty1, $name2 $ty2; $($name $ty),*);
  );
  (__impl $($name:ident $ty: ident),+; $name1:ident $ty1:ident, $($name2:ident $ty2:ident),*) => (
    permutation_trait_impl!($($name $ty),+);
    permutation_trait!(__impl $($name $ty),+ , $name1 $ty1; $($name2 $ty2),*);
  );
  (__impl $($name:ident $ty: ident),+; $name1:ident $ty1:ident) => (
    permutation_trait_impl!($($name $ty),+);
    permutation_trait_impl!($($name $ty),+, $name1 $ty1);
  );
);

macro_rules! permutation_trait_impl(
  ($($name:ident $ty: ident),+) => (
    impl<
      Input: Clone, $($ty),+ , Error: ParseError<Input>,
      $($name: Fn(Input) -> IResult<Input, $ty, Error>),+
    > Permutation<Input, ( $($ty),+ ), Error> for ( $($name),+ ) {

      fn permutation(&self, mut input: Input) -> IResult<Input, ( $($ty),+ ), Error> {
        let mut res = permutation_init!((), $($name),+);

        loop {
          let mut all_done = true;
          permutation_trait_inner!(0, self, input, res, all_done, $($name)+);

          //if we reach that part, it means none of the parsers were able to read anything
          if !all_done {
            //FIXME: should wrap the error returned by the child parser
            return Err(Err::Error(error_position!(input, ErrorKind::Permutation)));
          }
          break;
        }

        if let Some(unwrapped_res) = { permutation_trait_unwrap!(0, (), res, $($name),+) } {
          Ok((input, unwrapped_res))
        } else {
          Err(Err::Error(error_position!(input, ErrorKind::Permutation)))
        }
      }
    }
  );
);

macro_rules! permutation_trait_inner(
  ($it:tt, $self:expr, $input:ident, $res:expr, $all_done:expr, $head:ident $($id:ident)+) => ({
    if $res.$it.is_none() {
      match $self.$it($input.clone()) {
        Ok((i,o))     => {
          $input = i;
          $res.$it = Some(o);
          continue;
        },
        Err(Err::Error(_)) => {
          $all_done = false;
        },
        Err(e) => {
          return Err(e);
        }
      };
    }
    succ!($it, permutation_trait_inner!($self, $input, $res, $all_done, $($id)+));
  });
  ($it:tt, $self:expr, $input:ident, $res:expr, $all_done:expr, $head:ident) => ({
    if $res.$it.is_none() {
      match $self.$it($input.clone()) {
        Ok((i,o))     => {
          $input = i;
          $res.$it = Some(o);
          continue;
        },
        Err(Err::Error(_)) => {
          $all_done = false;
        },
        Err(e) => {
          return Err(e);
        }
      };
    }
  });
);

macro_rules! permutation_trait_unwrap (
  ($it:tt,  (), $res:ident, $e:ident, $($name:ident),+) => ({
    let res = $res.$it;
    if res.is_some() {
      succ!($it, permutation_trait_unwrap!((res.unwrap()), $res, $($name),+))
    } else {
      $crate::lib::std::option::Option::None
    }
  });

  ($it:tt, ($($parsed:expr),*), $res:ident, $e:ident, $($name:ident),+) => ({
    let res = $res.$it;
    if res.is_some() {
      succ!($it, permutation_trait_unwrap!(($($parsed),* , res.unwrap()), $res, $($name),+))
    } else {
      $crate::lib::std::option::Option::None
    }
  });

  ($it:tt, ($($parsed:expr),*), $res:ident, $name:ident) => ({
    let res = $res.$it;
    if res.is_some() {
      $crate::lib::std::option::Option::Some(($($parsed),* , res.unwrap() ))
    } else {
      $crate::lib::std::option::Option::None
    }
  });
);

permutation_trait!(FnA A, FnB B, FnC C, FnD D, FnE E, FnF F, FnG G, FnH H, FnI I, FnJ J, FnK K, FnL L, FnM M, FnN N, FnO O, FnP P, FnQ Q, FnR R, FnS S, FnT T, FnU U);
<|endoftext|>


<commit_before>use error::*;
use futures::future;
use futures::{Future, Stream};
use mime;
use multipart;
use hyper;
use hyper_tls::HttpsConnector;
use std::path::PathBuf;
use hyper::server;
use hyper::header::{Header, ContentDisposition, ContentType, DispositionParam, DispositionType};
use hyper::client::{Client, Request, Response};
use multipart::client::lazy;
use hyper::header::{Location};
use hyper::{Uri, StatusCode};
use std::io::prelude::*;
use tokio_core::reactor::Handle;

pub fn https_connector(handle: &Handle) -> HttpsConnector {
    HttpsConnector::new(4, handle)
}

pub trait ServerRequestExt {
    fn get_body_bytes(self) -> Box<Future<Item=Vec<u8>, Error=Error>>;
    fn has_content_type(&self, mime: mime::Mime) -> bool;
}

impl ServerRequestExt for server::Request {
    fn get_body_bytes(self) -> Box<Future<Item=Vec<u8>, Error=Error>> {
        Box::new(self.body().map_err(Error::from).fold(Vec::new(), |mut acc, chunk| {
            acc.extend_from_slice(&chunk);
            future::ok::<_, Error>(acc)
        }))
    }

    fn has_content_type(&self, mime: mime::Mime) -> bool {
        use mime::Mime;

        let content_type = self.headers().get::<ContentType>().cloned();
        let Mime(top_level, sub_level, _) = mime;

        if let Some(ContentType(Mime(found_top_level, found_sub_level, _))) = content_type {
            found_top_level == top_level && found_sub_level == sub_level
        } else {
            false
        }

    }
}

pub trait ResponseExt {
    fn file_name(&self) -> Option<String>;
    fn get_body_bytes(self) -> Box<Future<Item=Vec<u8>, Error=Error>>;
}

impl ResponseExt for Response {
    fn file_name(&self) -> Option<String> {
        match self.headers().get::<ContentDisposition>() {
            Some(&ContentDisposition { disposition: DispositionType::Attachment, parameters: ref params }) => {
                params.iter().find(|param| match **param {
                    DispositionParam::Filename(_, _, _) => true,
                    _ => false
                }).and_then(|param| {
                    if let DispositionParam::Filename(_, _, ref bytes) = *param {
                        String::from_utf8(bytes.to_owned()).ok()
                    } else {
                        None
                    }
                })
            }
            _ => None,
        }
    }

    fn get_body_bytes(self) -> Box<Future<Item=Vec<u8>, Error=Error>> {
        Box::new(self.body().map_err(Error::from).fold(Vec::new(), |mut acc, chunk| {
            acc.extend_from_slice(&chunk);
            future::ok::<_, Error>(acc)
        }))
    }
}

pub trait RequestExt {
    fn with_header<T: Header>(self, header: T) -> Self;

    fn with_body<T: Into<hyper::Body>>(self, body: T) -> Self;
}

impl RequestExt for Request<hyper::Body> {
    fn with_header<T: Header>(mut self, header: T) -> Self {
        {
            let mut h = self.headers_mut();
            h.set(header);
        }
        self
    }

    fn with_body<T: Into<hyper::Body>>(self, body: T) -> Self {
        let mut req = self;
        req.set_body(body.into());
        req
    }
}

pub trait ClientExt {
    fn get_follow_redirect(self, uri: &Uri) -> Box<Future<Item=Response, Error=Error>>;
}

impl ClientExt for Client<HttpsConnector> {
    fn get_follow_redirect(self, uri: &Uri) -> Box<Future<Item=Response, Error=Error>> {
        Box::new(future::loop_fn(uri.clone(), move |uri| {
            self.get(uri)
                .map_err(Error::from)
                .and_then(|res| {
                    match determine_get_result(res) {
                        Ok(GetResult::Redirect(redirect_uri)) => {
                            Ok(future::Loop::Continue(redirect_uri))
                        },
                        Ok(GetResult::Ok(res)) => Ok(future::Loop::Break(res)),
                        Err(err) => Err(err),
                    }
                })
        }))
    }
}

enum GetResult {
    Ok(Response),
    Redirect(Uri),
}

fn determine_get_result(res: Response) -> Result<GetResult> {
    match res.status() {
        StatusCode::TemporaryRedirect | StatusCode::PermanentRedirect => {
            match res.headers().get::<Location>() {
                Some(location) => Ok(GetResult::Redirect(location.parse()?)),
                None => Err(ErrorKind::UnprocessableEntity.into()),
            }
        },
        StatusCode::Ok => Ok(GetResult::Ok(res)),
        _ => Err(ErrorKind::UnprocessableEntity.into()),
    }
}

pub fn multipart_request_with_file(request: Request, path: PathBuf) -> ::std::result::Result<Request, Error> {
    let mut fields = lazy::Multipart::new()
        .add_file("file", path)
        .prepare_threshold(Some(u64::max_value() - 1))
        .map_err(|_| "Failed to prepare multipart body")?;
    let mut bytes: Vec<u8> = Vec::new();
    fields.read_to_end(&mut bytes)?;
    Ok(
        request
        .with_body(bytes)
        .with_header(ContentType(mime!(Multipart/FormData; Boundary=(fields.boundary()))))
    )
}

pub fn multipart_request_with_error(request: Request, error: &Error) -> Result<Request> {
    let mut fields = lazy::Multipart::new()
        .add_text("error", format!("{}", error))
        .prepare()
        .map_err(|_| "Failed to prepare multipart body")?;
    let mut bytes: Vec<u8> = Vec::new();
    fields.read_to_end(&mut bytes)?;
    Ok(
        request
        .with_body(bytes)
        .with_header(ContentType(mime!(Multipart/FormData; Boundary=(fields.boundary()))))
    )
}

#[derive(Debug)]
pub struct MultipartRequest(pub hyper::Headers, pub Vec<u8>);

impl multipart::server::HttpRequest for MultipartRequest {
    type Body = ::std::io::Cursor<Vec<u8>>;

    fn multipart_boundary(&self) -> Option<&str> {
        let content_type = self.0.get::<ContentType>();
        match content_type {
            Some(&ContentType(mime::Mime(mime::TopLevel::Multipart, mime::SubLevel::FormData, ref params))) => {
                // param is (attr, value)
                params.iter().find(|param| {
                    param.0.as_str() == "boundary"
                }).map(|param| param.1.as_str())
            },
            _ => None
        }
    }

    fn body(self) -> Self::Body {
        ::std::io::Cursor::new(self.1)
    }
}
<commit_msg>🐛 Avoid UnprocessableEntity in the client<commit_after>use error::*;
use futures::future;
use futures::{Future, Stream};
use mime;
use multipart;
use hyper;
use hyper_tls::HttpsConnector;
use std::path::PathBuf;
use hyper::server;
use hyper::header::{Header, ContentDisposition, ContentType, DispositionParam, DispositionType};
use hyper::client::{Client, Request, Response};
use multipart::client::lazy;
use hyper::header::{Location};
use hyper::{Uri, StatusCode};
use std::io::prelude::*;
use tokio_core::reactor::Handle;

pub fn https_connector(handle: &Handle) -> HttpsConnector {
    HttpsConnector::new(4, handle)
}

pub trait ServerRequestExt {
    fn get_body_bytes(self) -> Box<Future<Item=Vec<u8>, Error=Error>>;
    fn has_content_type(&self, mime: mime::Mime) -> bool;
}

impl ServerRequestExt for server::Request {
    fn get_body_bytes(self) -> Box<Future<Item=Vec<u8>, Error=Error>> {
        Box::new(self.body().map_err(Error::from).fold(Vec::new(), |mut acc, chunk| {
            acc.extend_from_slice(&chunk);
            future::ok::<_, Error>(acc)
        }))
    }

    fn has_content_type(&self, mime: mime::Mime) -> bool {
        use mime::Mime;

        let content_type = self.headers().get::<ContentType>().cloned();
        let Mime(top_level, sub_level, _) = mime;

        if let Some(ContentType(Mime(found_top_level, found_sub_level, _))) = content_type {
            found_top_level == top_level && found_sub_level == sub_level
        } else {
            false
        }

    }
}

pub trait ResponseExt {
    fn file_name(&self) -> Option<String>;
    fn get_body_bytes(self) -> Box<Future<Item=Vec<u8>, Error=Error>>;
}

impl ResponseExt for Response {
    fn file_name(&self) -> Option<String> {
        match self.headers().get::<ContentDisposition>() {
            Some(&ContentDisposition { disposition: DispositionType::Attachment, parameters: ref params }) => {
                params.iter().find(|param| match **param {
                    DispositionParam::Filename(_, _, _) => true,
                    _ => false
                }).and_then(|param| {
                    if let DispositionParam::Filename(_, _, ref bytes) = *param {
                        String::from_utf8(bytes.to_owned()).ok()
                    } else {
                        None
                    }
                })
            }
            _ => None,
        }
    }

    fn get_body_bytes(self) -> Box<Future<Item=Vec<u8>, Error=Error>> {
        Box::new(self.body().map_err(Error::from).fold(Vec::new(), |mut acc, chunk| {
            acc.extend_from_slice(&chunk);
            future::ok::<_, Error>(acc)
        }))
    }
}

pub trait RequestExt {
    fn with_header<T: Header>(self, header: T) -> Self;

    fn with_body<T: Into<hyper::Body>>(self, body: T) -> Self;
}

impl RequestExt for Request<hyper::Body> {
    fn with_header<T: Header>(mut self, header: T) -> Self {
        {
            let mut h = self.headers_mut();
            h.set(header);
        }
        self
    }

    fn with_body<T: Into<hyper::Body>>(self, body: T) -> Self {
        let mut req = self;
        req.set_body(body.into());
        req
    }
}

pub trait ClientExt {
    fn get_follow_redirect(self, uri: &Uri) -> Box<Future<Item=Response, Error=Error>>;
}

impl ClientExt for Client<HttpsConnector> {
    fn get_follow_redirect(self, uri: &Uri) -> Box<Future<Item=Response, Error=Error>> {
        Box::new(future::loop_fn(uri.clone(), move |uri| {
            self.get(uri)
                .map_err(Error::from)
                .and_then(|res| {
                    match determine_get_result(res) {
                        Ok(GetResult::Redirect(redirect_uri)) => {
                            Ok(future::Loop::Continue(redirect_uri))
                        },
                        Ok(GetResult::Ok(res)) => Ok(future::Loop::Break(res)),
                        Err(err) => Err(err),
                    }
                })
        }))
    }
}

enum GetResult {
    Ok(Response),
    Redirect(Uri),
}

fn determine_get_result(res: Response) -> Result<GetResult> {
    match res.status() {
        StatusCode::TemporaryRedirect | StatusCode::PermanentRedirect => {
            match res.headers().get::<Location>() {
                Some(location) => Ok(GetResult::Redirect(location.parse()?)),
                None => Err("Redirect without Location header".into()),
            }
        },
        _ => Ok(GetResult::Ok(res)),
    }
}

pub fn multipart_request_with_file(request: Request, path: PathBuf) -> ::std::result::Result<Request, Error> {
    let mut fields = lazy::Multipart::new()
        .add_file("file", path)
        .prepare_threshold(Some(u64::max_value() - 1))
        .map_err(|_| "Failed to prepare multipart body")?;
    let mut bytes: Vec<u8> = Vec::new();
    fields.read_to_end(&mut bytes)?;
    Ok(
        request
        .with_body(bytes)
        .with_header(ContentType(mime!(Multipart/FormData; Boundary=(fields.boundary()))))
    )
}

pub fn multipart_request_with_error(request: Request, error: &Error) -> Result<Request> {
    let mut fields = lazy::Multipart::new()
        .add_text("error", format!("{}", error))
        .prepare()
        .map_err(|_| "Failed to prepare multipart body")?;
    let mut bytes: Vec<u8> = Vec::new();
    fields.read_to_end(&mut bytes)?;
    Ok(
        request
        .with_body(bytes)
        .with_header(ContentType(mime!(Multipart/FormData; Boundary=(fields.boundary()))))
    )
}

#[derive(Debug)]
pub struct MultipartRequest(pub hyper::Headers, pub Vec<u8>);

impl multipart::server::HttpRequest for MultipartRequest {
    type Body = ::std::io::Cursor<Vec<u8>>;

    fn multipart_boundary(&self) -> Option<&str> {
        let content_type = self.0.get::<ContentType>();
        match content_type {
            Some(&ContentType(mime::Mime(mime::TopLevel::Multipart, mime::SubLevel::FormData, ref params))) => {
                // param is (attr, value)
                params.iter().find(|param| {
                    param.0.as_str() == "boundary"
                }).map(|param| param.1.as_str())
            },
            _ => None
        }
    }

    fn body(self) -> Self::Body {
        ::std::io::Cursor::new(self.1)
    }
}
<|endoftext|>


<commit_before>// Copyright 2014 Alexis Mousset. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! SMTP client

use std::fmt::Show;
use std::result::Result;
use std::string::String;
use std::io::net::ip::Port;

use common::{get_first_word, unquote_email_address};
use response::Response;
use extension;
use command;
use command::Command;
use common::{SMTP_PORT, CRLF};
use transaction;
use transaction::TransactionState;
use client::connecter::Connecter;
use client::server_info::ServerInfo;
use client::stream::ClientStream;

mod connecter;
mod server_info;
mod stream;

/// Structure that implements the SMTP client
pub struct Client<S> {
    /// TCP stream between client and server
    /// Value is None before connection
    stream: Option<S>,
    /// Host we are connecting to
    host: String,
    /// Port we are connecting on
    port: Port,
    /// Our hostname for HELO/EHLO commands
    my_hostname: String,
    /// Information about the server
    /// Value is None before HELO/EHLO
    server_info: Option<ServerInfo>,
    /// Transaction state, to check the sequence of commands
    state: TransactionState
}

impl<S> Client<S> {
    /// Creates a new SMTP client
    pub fn new(host: String, port: Option<Port>, my_hostname: Option<String>) -> Client<S> {
        Client{
            stream: None,
            host: host,
            port: port.unwrap_or(SMTP_PORT),
            my_hostname: my_hostname.unwrap_or(String::from_str("localhost")),
            server_info: None,
            state: transaction::Unconnected
        }
    }
}

// T : String ou String, selon le support
impl<S: Connecter + ClientStream + Clone> Client<S> {

    /// TODO
    fn smtp_fail_if_err<S>(&mut self, response: Result<Response, Response>) {
        match response {
            Err(response) => {
                self.smtp_fail::<S, Response>(response)
            },
            Ok(_) => {}
        }
    }

    /// Connects to the configured server
    pub fn connect(&mut self) -> Result<Response, Response> {
        // connect should not be called when the client is already connected
        if !self.stream.is_none() {
            panic!("The connection is already established");
        }

        // Try to connect
        self.stream = match Connecter::connect(self.host.clone().as_slice(), self.port) {
            Ok(stream) => Some(stream),
            Err(..) => panic!("Cannot connect to the server")
        };

        // Log the connection
        info!("Connection established to {}[{}]:{}",
              self.host, self.stream.clone().unwrap().peer_name().unwrap().ip, self.port);

        match self.stream.clone().unwrap().get_reply() {
            Some(response) => match response.with_code(vec!(220)) {
                                  Ok(response)  => {
                                      self.state = transaction::Connected;
                                      Ok(response)
                                  },
                                  Err(response) => {
                                      Err(response)
                                  }
                              },
            None           => panic!("No banner on {}", self.host)
        }
    }

    /// Sends an email
    pub fn send_mail<S>(&mut self, from_address: String,
                        to_addresses: Vec<String>, message: String) {
        let my_hostname = self.my_hostname.clone();
        let mut smtp_result: Result<Response, Response>;

        match self.connect() {
            Ok(_) => {},
            Err(response) => panic!("Cannot connect to {:s}:{:u}. Server says: {}",
                                    self.host,
                                    self.port, response
                             )
        }

        // Extended Hello or Hello
        match self.ehlo::<S>(my_hostname.clone().to_string()) {
            Err(Response{code: 550, message: _}) => {
                smtp_result = self.helo::<S>(my_hostname.clone());
                self.smtp_fail_if_err::<S>(smtp_result);
            },
            Err(response) => {
                self.smtp_fail::<S, Response>(response)
            }
            _ => {}
        }

        debug!("Server {}", self.server_info.clone().unwrap());

        // Checks message encoding according to the server's capability
        // TODO : Add an encoding check.
        if ! self.server_info.clone().unwrap().supports_feature(extension::EightBitMime).is_ok() {
            if ! message.clone().to_string().is_ascii() {
                self.smtp_fail::<S, &str>("Server does not accepts UTF-8 strings");
            }
        }

        // Mail
        smtp_result = self.mail::<S>(from_address.clone(), None);
        self.smtp_fail_if_err::<S>(smtp_result);

        // Log the mail command
        info!("from=<{}>, size={}, nrcpt={}", from_address, 42u, to_addresses.len());

        // Recipient
        // TODO Return rejected addresses
        // TODO Manage the number of recipients
        for to_address in to_addresses.iter() {
            smtp_result = self.rcpt::<S>(to_address.clone(), None);
            self.smtp_fail_if_err::<S>(smtp_result);
        }

        // Data
        smtp_result = self.data::<S>();
        self.smtp_fail_if_err::<S>(smtp_result);

        // Message content
        let sent = self.message::<S>(message);

        if sent.clone().is_err() {
            self.smtp_fail::<S, Response>(sent.clone().err().unwrap())
        }

        info!("to=<{}>, status=sent ({})",
              to_addresses.clone().connect(">, to=<"), sent.clone().ok().unwrap());

        // Quit
        smtp_result = self.quit::<S>();
        self.smtp_fail_if_err::<S>(smtp_result);
    }

    /// Sends an SMTP command
    // TODO : ensure this is an ASCII string
    fn send_command(&mut self, command: Command) -> Response {
        if !self.state.is_command_possible(command.clone()) {
            panic!("Bad command sequence");
        }
        self.stream.clone().unwrap().send_and_get_response(format!("{}", command).as_slice())
    }

    /// Sends an email
    fn send_message(&mut self, message: String) -> Response {
        self.stream.clone().unwrap().send_and_get_response(format!("{}{}.{}", message, CRLF, CRLF).as_slice())
    }

    /// Closes the connection and fail with a given messgage
    fn smtp_fail<S, T: Show>(&mut self, reason: T) {
        let is_connected = self.is_connected::<S>();
        if is_connected {
            match self.quit::<S>() {
                Ok(..) => {},
                Err(response) => panic!("Failed: {}", response)
            }
        }
        self.close();
        panic!("Failed: {}", reason);
    }

    /// Checks if the server is connected
    pub fn is_connected<S>(&mut self) -> bool {
        self.noop::<S>().is_ok()
    }

    /// Closes the TCP stream
    pub fn close(&mut self) {
        // Close the TCP connection
        drop(self.stream.clone().unwrap());
        // Reset client state
        self.stream = None;
        self.state = transaction::Unconnected;
        self.server_info = None;
    }

    /// Send a HELO command
    pub fn helo<S>(&mut self, my_hostname: String) -> Result<Response, Response> {
        match self.send_command(command::Hello(my_hostname.clone())).with_code(vec!(250)) {
            Ok(response) => {
                self.server_info = Some(
                    ServerInfo{
                        name: get_first_word(response.message.clone().unwrap()),
                        esmtp_features: None
                    }
                );
                self.state = transaction::HelloSent;
                Ok(response)
            },
            Err(response) => Err(response)
        }
    }

    /// Sends a EHLO command
    pub fn ehlo<S>(&mut self, my_hostname: String) -> Result<Response, Response> {
        match self.send_command(command::ExtendedHello(my_hostname.clone())).with_code(vec!(250)) {
            Ok(response) => {
                self.server_info = Some(
                    ServerInfo{
                        name: get_first_word(response.message.clone().unwrap()),
                        esmtp_features: ServerInfo::parse_esmtp_response(
                                            response.message.clone().unwrap()
                                        )
                    }
                );
                self.state = transaction::HelloSent;
                Ok(response)
            },
            Err(response) => Err(response)
        }
    }

    /// Sends a MAIL command
    pub fn mail<S>(&mut self, from_address: String,
                   options: Option<Vec<String>>) -> Result<Response, Response> {
        match self.send_command(
            command::Mail(unquote_email_address(from_address.to_string()), options)
        ).with_code(vec!(250)) {
            Ok(response) => {
                self.state = transaction::MailSent;
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a RCPT command
    pub fn rcpt<S>(&mut self, to_address: String,
                   options: Option<Vec<String>>) -> Result<Response, Response> {
        match self.send_command(
            command::Recipient(unquote_email_address(to_address.to_string()), options)
        ).with_code(vec!(250)) {
            Ok(response) => {
                self.state = transaction::RecipientSent;
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a DATA command
    pub fn data<S>(&mut self) -> Result<Response, Response> {
        match self.send_command(command::Data).with_code(vec!(354)) {
            Ok(response) => {
                self.state = transaction::DataSent;
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends the message content
    pub fn message<S>(&mut self, message_content: String) -> Result<Response, Response> {
        match self.send_message(message_content).with_code(vec!(250)) {
            Ok(response) => {
                self.state = transaction::HelloSent;
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a QUIT command
    pub fn quit<S>(&mut self) -> Result<Response, Response> {
        match self.send_command(command::Quit).with_code(vec!(221)) {
            Ok(response) => {
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a RSET command
    pub fn rset<S>(&mut self) -> Result<Response, Response> {
        match self.send_command(command::Reset).with_code(vec!(250)) {
            Ok(response) => {
                if vec!(transaction::MailSent, transaction::RecipientSent,
                        transaction::DataSent).contains(&self.state) {
                    self.state = transaction::HelloSent;
                }
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a NOOP commands
    pub fn noop<S>(&mut self) -> Result<Response, Response> {
        self.send_command(command::Noop).with_code(vec!(250))
    }

    /// Sends a VRFY command
    pub fn vrfy<S, T>(&mut self, to_address: String) -> Result<Response, Response> {
        self.send_command(command::Verify(to_address, None)).with_code(vec!(250))
    }
}
<commit_msg>Use the extensions supported by the server<commit_after>// Copyright 2014 Alexis Mousset. See the COPYRIGHT
// file at the top-level directory of this distribution.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

//! SMTP client

use std::fmt::Show;
use std::result::Result;
use std::string::String;
use std::io::net::ip::Port;

use common::{get_first_word, unquote_email_address};
use response::Response;
use extension;
use command;
use command::Command;
use common::{SMTP_PORT, CRLF};
use transaction;
use transaction::TransactionState;
use client::connecter::Connecter;
use client::server_info::ServerInfo;
use client::stream::ClientStream;

mod connecter;
mod server_info;
mod stream;

/// Structure that implements the SMTP client
pub struct Client<S> {
    /// TCP stream between client and server
    /// Value is None before connection
    stream: Option<S>,
    /// Host we are connecting to
    host: String,
    /// Port we are connecting on
    port: Port,
    /// Our hostname for HELO/EHLO commands
    my_hostname: String,
    /// Information about the server
    /// Value is None before HELO/EHLO
    server_info: Option<ServerInfo>,
    /// Transaction state, to check the sequence of commands
    state: TransactionState
}

impl<S> Client<S> {
    /// Creates a new SMTP client
    pub fn new(host: String, port: Option<Port>, my_hostname: Option<String>) -> Client<S> {
        Client{
            stream: None,
            host: host,
            port: port.unwrap_or(SMTP_PORT),
            my_hostname: my_hostname.unwrap_or(String::from_str("localhost")),
            server_info: None,
            state: transaction::Unconnected
        }
    }
}

// T : String ou String, selon le support
impl<S: Connecter + ClientStream + Clone> Client<S> {

    /// TODO
    fn smtp_fail_if_err<S>(&mut self, response: Result<Response, Response>) {
        match response {
            Err(response) => {
                self.smtp_fail::<S, Response>(response)
            },
            Ok(_) => {}
        }
    }

    /// Closes the connection and fail with a given messgage
    fn smtp_fail<S, T: Show>(&mut self, reason: T) {
        let is_connected = self.is_connected::<S>();
        if is_connected {
            match self.quit::<S>() {
                Ok(..) => {},
                Err(response) => panic!("Failed: {}", response)
            }
        }
        self.close();
        panic!("Failed: {}", reason);
    }

    /// Sends an email
    pub fn send_mail<S>(&mut self, from_address: String,
                        to_addresses: Vec<String>, message: String) {
        let my_hostname = self.my_hostname.clone();
        let mut smtp_result: Result<Response, Response>;

        match self.connect() {
            Ok(_) => {},
            Err(response) => panic!("Cannot connect to {:s}:{:u}. Server says: {}",
                                    self.host,
                                    self.port, response
                             )
        }

        // Extended Hello or Hello
        match self.ehlo::<S>(my_hostname.clone().to_string()) {
            Err(Response{code: 550, message: _}) => {
                smtp_result = self.helo::<S>(my_hostname.clone());
                self.smtp_fail_if_err::<S>(smtp_result);
            },
            Err(response) => {
                self.smtp_fail::<S, Response>(response)
            }
            _ => {}
        }

        debug!("Server {}", self.server_info.clone().unwrap());

        // Checks message encoding according to the server's capability
        // TODO : Add an encoding check.
        if ! self.server_info.clone().unwrap().supports_feature(extension::EightBitMime).is_ok() {
            if ! message.clone().to_string().is_ascii() {
                self.smtp_fail::<S, &str>("Server does not accepts UTF-8 strings");
            }
        }

        // Get maximum message size if defined
        let max_size = match self.server_info.clone().unwrap().supports_feature(extension::Size(0)) {
            Ok(extension::Size(max)) => max,
            _ => -1
        };

        // Check maximum message size
        if max_size > 0 && message.len() > max_size {
            self.smtp_fail::<S, &str>("Message is too big. The limit is {}");
        }

        // Mail
        smtp_result = self.mail::<S>(from_address.clone(), None);
        self.smtp_fail_if_err::<S>(smtp_result);

        // Log the mail command
        info!("from=<{}>, size={}, nrcpt={}", from_address, message.len(), to_addresses.len());

        // Recipient
        // TODO Return rejected addresses
        // TODO Manage the number of recipients
        for to_address in to_addresses.iter() {
            smtp_result = self.rcpt::<S>(to_address.clone(), None);
            self.smtp_fail_if_err::<S>(smtp_result);
        }

        // Data
        smtp_result = self.data::<S>();
        self.smtp_fail_if_err::<S>(smtp_result);

        // Message content
        let sent = self.message::<S>(message);

        if sent.clone().is_err() {
            self.smtp_fail::<S, Response>(sent.clone().err().unwrap())
        }

        info!("to=<{}>, status=sent ({})",
              to_addresses.clone().connect(">, to=<"), sent.clone().ok().unwrap());

        // Quit
        smtp_result = self.quit::<S>();
        self.smtp_fail_if_err::<S>(smtp_result);
    }

    /// Sends an SMTP command
    // TODO : ensure this is an ASCII string
    fn send_command(&mut self, command: Command) -> Response {
        if !self.state.is_command_possible(command.clone()) {
            panic!("Bad command sequence");
        }
        self.stream.clone().unwrap().send_and_get_response(format!("{}", command).as_slice())
    }

    /// Sends the email content
    fn send_message(&mut self, message: String) -> Response {
        self.stream.clone().unwrap().send_and_get_response(format!("{}{}.{}", message, CRLF, CRLF).as_slice())
    }

    /// Connects to the configured server
    pub fn connect(&mut self) -> Result<Response, Response> {
        // connect should not be called when the client is already connected
        if !self.stream.is_none() {
            panic!("The connection is already established");
        }

        // Try to connect
        self.stream = match Connecter::connect(self.host.clone().as_slice(), self.port) {
            Ok(stream) => Some(stream),
            Err(..) => panic!("Cannot connect to the server")
        };

        // Log the connection
        info!("Connection established to {}[{}]:{}",
              self.host, self.stream.clone().unwrap().peer_name().unwrap().ip, self.port);

        match self.stream.clone().unwrap().get_reply() {
            Some(response) => match response.with_code(vec!(220)) {
                                  Ok(response)  => {
                                      self.state = transaction::Connected;
                                      Ok(response)
                                  },
                                  Err(response) => {
                                      Err(response)
                                  }
                              },
            None           => panic!("No banner on {}", self.host)
        }
    }

    /// Checks if the server is connected
    pub fn is_connected<S>(&mut self) -> bool {
        self.noop::<S>().is_ok()
    }

    /// Closes the TCP stream
    pub fn close(&mut self) {
        // Close the TCP connection
        drop(self.stream.clone().unwrap());
        // Reset client state
        self.stream = None;
        self.state = transaction::Unconnected;
        self.server_info = None;
    }

    /// Send a HELO command
    pub fn helo<S>(&mut self, my_hostname: String) -> Result<Response, Response> {
        match self.send_command(command::Hello(my_hostname.clone())).with_code(vec!(250)) {
            Ok(response) => {
                self.server_info = Some(
                    ServerInfo{
                        name: get_first_word(response.message.clone().unwrap()),
                        esmtp_features: None
                    }
                );
                self.state = transaction::HelloSent;
                Ok(response)
            },
            Err(response) => Err(response)
        }
    }

    /// Sends a EHLO command
    pub fn ehlo<S>(&mut self, my_hostname: String) -> Result<Response, Response> {
        match self.send_command(command::ExtendedHello(my_hostname.clone())).with_code(vec!(250)) {
            Ok(response) => {
                self.server_info = Some(
                    ServerInfo{
                        name: get_first_word(response.message.clone().unwrap()),
                        esmtp_features: ServerInfo::parse_esmtp_response(
                                            response.message.clone().unwrap()
                                        )
                    }
                );
                self.state = transaction::HelloSent;
                Ok(response)
            },
            Err(response) => Err(response)
        }
    }

    /// Sends a MAIL command
    pub fn mail<S>(&mut self, from_address: String,
                   options: Option<Vec<String>>) -> Result<Response, Response> {
        match self.send_command(
            command::Mail(unquote_email_address(from_address.to_string()), options)
        ).with_code(vec!(250)) {
            Ok(response) => {
                self.state = transaction::MailSent;
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a RCPT command
    pub fn rcpt<S>(&mut self, to_address: String,
                   options: Option<Vec<String>>) -> Result<Response, Response> {
        match self.send_command(
            command::Recipient(unquote_email_address(to_address.to_string()), options)
        ).with_code(vec!(250)) {
            Ok(response) => {
                self.state = transaction::RecipientSent;
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a DATA command
    pub fn data<S>(&mut self) -> Result<Response, Response> {
        match self.send_command(command::Data).with_code(vec!(354)) {
            Ok(response) => {
                self.state = transaction::DataSent;
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends the message content
    pub fn message<S>(&mut self, message_content: String) -> Result<Response, Response> {
        match self.send_message(message_content).with_code(vec!(250)) {
            Ok(response) => {
                self.state = transaction::HelloSent;
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a QUIT command
    pub fn quit<S>(&mut self) -> Result<Response, Response> {
        match self.send_command(command::Quit).with_code(vec!(221)) {
            Ok(response) => {
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a RSET command
    pub fn rset<S>(&mut self) -> Result<Response, Response> {
        match self.send_command(command::Reset).with_code(vec!(250)) {
            Ok(response) => {
                if vec!(transaction::MailSent, transaction::RecipientSent,
                        transaction::DataSent).contains(&self.state) {
                    self.state = transaction::HelloSent;
                }
                Ok(response)
            },
            Err(response) => {
                Err(response)
            }
        }
    }

    /// Sends a NOOP commands
    pub fn noop<S>(&mut self) -> Result<Response, Response> {
        self.send_command(command::Noop).with_code(vec!(250))
    }

    /// Sends a VRFY command
    pub fn vrfy<S, T>(&mut self, to_address: String) -> Result<Response, Response> {
        self.send_command(command::Verify(to_address, None)).with_code(vec!(250))
    }
}
<|endoftext|>


<commit_before>///! An item is line of text that read from `find` command or stdin together with
///! the internal states, such as selected or not
use crate::ansi::{ANSIParser, AnsiString};
use crate::field::*;
use regex::Regex;
use std::borrow::Cow;
use std::cmp::Ordering;
use std::default::Default;
use std::sync::Arc;
use crate::spinlock::{SpinLock, SpinLockGuard};
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use std::iter::Iterator;

/// An item will store everything that one line input will need to be operated and displayed.
///
/// What's special about an item?
/// The simplest version of an item is a line of string, but things are getting more complex:
/// - The conversion of lower/upper case is slow in rust, because it involds unicode.
/// - We may need to interpret the ANSI codes in the text.
/// - The text can be transformed and limited while searching.
///
/// About the ANSI, we made assumption that it is linewise, that means no ANSI codes will affect
/// more than one line.
#[derive(Debug)]
pub struct Item {
    // (num of run, number of index)
    index: (usize, usize),

    // The text that will be ouptut when user press `enter`
    orig_text: String,

    // The text that will shown into the screen. Can be transformed.
    text: AnsiString,

    matching_ranges: Vec<(usize, usize)>,

    // For the transformed ANSI case, the output will need another transform.
    using_transform_fields: bool,
    ansi_enabled: bool,
}

impl<'a> Item {
    pub fn new(
        orig_text: Cow<str>,
        ansi_enabled: bool,
        trans_fields: &[FieldRange],
        matching_fields: &[FieldRange],
        delimiter: &Regex,
        index: (usize, usize),
    ) -> Self {
        let using_transform_fields = !trans_fields.is_empty();

        //        transformed | ANSI             | output
        //------------------------------------------------------
        //                    +- T -> trans+ANSI | ANSI
        //                    |                  |
        //      +- T -> trans +- F -> trans      | orig
        // orig |                                |
        //      +- F -> orig  +- T -> ANSI     ==| ANSI
        //                    |                  |
        //                    +- F -> orig       | orig

        let mut ansi_parser: ANSIParser = Default::default();

        let text = if using_transform_fields && ansi_enabled {
            // ansi and transform
            ansi_parser.parse_ansi(&parse_transform_fields(delimiter, &orig_text, trans_fields))
        } else if using_transform_fields {
            // transformed, not ansi
            AnsiString::new_string(parse_transform_fields(delimiter, &orig_text, trans_fields))
        } else if ansi_enabled {
            // not transformed, ansi
            ansi_parser.parse_ansi(&orig_text)
        } else {
            // normal case
            AnsiString::new_empty()
        };

        let mut ret = Item {
            index,
            orig_text: orig_text.into_owned(),
            text,
            using_transform_fields: !trans_fields.is_empty(),
            matching_ranges: Vec::new(),
            ansi_enabled,
        };

        let matching_ranges = if !matching_fields.is_empty() {
            parse_matching_fields(delimiter, ret.get_text(), matching_fields)
        } else {
            vec![(0, ret.get_text().len())]
        };

        ret.matching_ranges = matching_ranges;
        ret
    }

    pub fn get_text(&self) -> &str {
        if !self.using_transform_fields && !self.ansi_enabled {
            &self.orig_text
        } else {
            &self.text.get_stripped()
        }
    }

    pub fn get_text_struct(&self) -> Option<&AnsiString> {
        if !self.using_transform_fields && !self.ansi_enabled {
            None
        } else {
            Some(&self.text)
        }
    }

    pub fn get_output_text(&'a self) -> Cow<'a, str> {
        if self.using_transform_fields && self.ansi_enabled {
            let mut ansi_parser: ANSIParser = Default::default();
            let text = ansi_parser.parse_ansi(&self.orig_text);
            Cow::Owned(text.into_inner())
        } else if !self.using_transform_fields && self.ansi_enabled {
            Cow::Borrowed(self.text.get_stripped())
        } else {
            Cow::Borrowed(&self.orig_text)
        }
    }

    pub fn get_index(&self) -> usize {
        self.index.1
    }

    pub fn get_full_index(&self) -> (usize, usize) {
        self.index
    }

    pub fn get_matching_ranges(&self) -> &[(usize, usize)] {
        &self.matching_ranges
    }
}

impl Clone for Item {
    fn clone(&self) -> Item {
        Item {
            index: self.index,
            orig_text: self.orig_text.clone(),
            text: self.text.clone(),
            using_transform_fields: self.using_transform_fields,
            matching_ranges: self.matching_ranges.clone(),
            ansi_enabled: self.ansi_enabled,
        }
    }
}

pub type Rank = [i64; 4]; // score, index, start, end

#[derive(PartialEq, Eq, Clone, Debug)]
#[allow(dead_code)]
pub enum MatchedRange {
    ByteRange(usize, usize), // range of bytes
    Chars(Vec<usize>),       // individual characters matched
}

#[derive(Clone, Debug)]
pub struct MatchedItem {
    pub item: Arc<Item>,
    pub rank: Rank,
    pub matched_range: Option<MatchedRange>, // range of chars that matched the pattern
}

impl MatchedItem {
    pub fn builder(item: Arc<Item>) -> Self {
        MatchedItem {
            item,
            rank: [0, 0, 0, 0],
            matched_range: None,
        }
    }

    pub fn matched_range(mut self, range: MatchedRange) -> Self {
        self.matched_range = Some(range);
        self
    }

    pub fn rank(mut self, rank: Rank) -> Self {
        self.rank = rank;
        self
    }

    pub fn build(self) -> Self {
        self
    }
}

impl Ord for MatchedItem {
    fn cmp(&self, other: &MatchedItem) -> Ordering {
        self.rank.cmp(&other.rank)
    }
}

// `PartialOrd` needs to be implemented as well.
impl PartialOrd for MatchedItem {
    fn partial_cmp(&self, other: &MatchedItem) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for MatchedItem {
    fn eq(&self, other: &MatchedItem) -> bool {
        self.rank == other.rank
    }
}

impl Eq for MatchedItem {}

const ITEM_POOL_CAPACITY : usize = 1024;

pub struct ItemPool {
    pool: SpinLock<Vec<Arc<Item>>>,
    /// number of items that was `take`n
    taken: AtomicUsize,
}

impl ItemPool {
    pub fn new() -> Self {
        Self {
            pool: SpinLock::new(Vec::with_capacity(ITEM_POOL_CAPACITY)),
            taken: AtomicUsize::new(0),
        }
    }

    pub fn len(&self) -> usize {
        let items = self.pool.lock();
        items.len()
    }

    pub fn clear(&self) {
        let mut items = self.pool.lock();
        items.clear();
        self.taken.store(0, AtomicOrdering::SeqCst);
    }

    pub fn reset(&self) {
        // lock to ensure consistency
        let items = self.pool.lock();
        self.taken.store(0, AtomicOrdering::SeqCst);
    }

    pub fn append(&self, items: &mut Vec<Arc<Item>>) {
        let mut pool = self.pool.lock();
        pool.append(items);
    }

    pub fn take(&self) -> Vec<Arc<Item>> {
        let pool = self.pool.lock();
        let len = pool.len();
        let taken = self.taken.swap(len, AtomicOrdering::SeqCst);
        let mut ret = Vec::with_capacity(len-taken);
        for item in &pool[taken..len] {
            ret.push(item.clone())
        }
        ret
    }
}
<commit_msg>[item] avoid allocate on take for item pool<commit_after>///! An item is line of text that read from `find` command or stdin together with
///! the internal states, such as selected or not
use crate::ansi::{ANSIParser, AnsiString};
use crate::field::*;
use regex::Regex;
use std::borrow::Cow;
use std::cmp::Ordering;
use std::default::Default;
use std::sync::Arc;
use crate::spinlock::{SpinLock, SpinLockGuard};
use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering};
use std::iter::Iterator;
use std::ops::Deref;

/// An item will store everything that one line input will need to be operated and displayed.
///
/// What's special about an item?
/// The simplest version of an item is a line of string, but things are getting more complex:
/// - The conversion of lower/upper case is slow in rust, because it involds unicode.
/// - We may need to interpret the ANSI codes in the text.
/// - The text can be transformed and limited while searching.
///
/// About the ANSI, we made assumption that it is linewise, that means no ANSI codes will affect
/// more than one line.
#[derive(Debug)]
pub struct Item {
    // (num of run, number of index)
    index: (usize, usize),

    // The text that will be ouptut when user press `enter`
    orig_text: String,

    // The text that will shown into the screen. Can be transformed.
    text: AnsiString,

    matching_ranges: Vec<(usize, usize)>,

    // For the transformed ANSI case, the output will need another transform.
    using_transform_fields: bool,
    ansi_enabled: bool,
}

impl<'a> Item {
    pub fn new(
        orig_text: Cow<str>,
        ansi_enabled: bool,
        trans_fields: &[FieldRange],
        matching_fields: &[FieldRange],
        delimiter: &Regex,
        index: (usize, usize),
    ) -> Self {
        let using_transform_fields = !trans_fields.is_empty();

        //        transformed | ANSI             | output
        //------------------------------------------------------
        //                    +- T -> trans+ANSI | ANSI
        //                    |                  |
        //      +- T -> trans +- F -> trans      | orig
        // orig |                                |
        //      +- F -> orig  +- T -> ANSI     ==| ANSI
        //                    |                  |
        //                    +- F -> orig       | orig

        let mut ansi_parser: ANSIParser = Default::default();

        let text = if using_transform_fields && ansi_enabled {
            // ansi and transform
            ansi_parser.parse_ansi(&parse_transform_fields(delimiter, &orig_text, trans_fields))
        } else if using_transform_fields {
            // transformed, not ansi
            AnsiString::new_string(parse_transform_fields(delimiter, &orig_text, trans_fields))
        } else if ansi_enabled {
            // not transformed, ansi
            ansi_parser.parse_ansi(&orig_text)
        } else {
            // normal case
            AnsiString::new_empty()
        };

        let mut ret = Item {
            index,
            orig_text: orig_text.into_owned(),
            text,
            using_transform_fields: !trans_fields.is_empty(),
            matching_ranges: Vec::new(),
            ansi_enabled,
        };

        let matching_ranges = if !matching_fields.is_empty() {
            parse_matching_fields(delimiter, ret.get_text(), matching_fields)
        } else {
            vec![(0, ret.get_text().len())]
        };

        ret.matching_ranges = matching_ranges;
        ret
    }

    pub fn get_text(&self) -> &str {
        if !self.using_transform_fields && !self.ansi_enabled {
            &self.orig_text
        } else {
            &self.text.get_stripped()
        }
    }

    pub fn get_text_struct(&self) -> Option<&AnsiString> {
        if !self.using_transform_fields && !self.ansi_enabled {
            None
        } else {
            Some(&self.text)
        }
    }

    pub fn get_output_text(&'a self) -> Cow<'a, str> {
        if self.using_transform_fields && self.ansi_enabled {
            let mut ansi_parser: ANSIParser = Default::default();
            let text = ansi_parser.parse_ansi(&self.orig_text);
            Cow::Owned(text.into_inner())
        } else if !self.using_transform_fields && self.ansi_enabled {
            Cow::Borrowed(self.text.get_stripped())
        } else {
            Cow::Borrowed(&self.orig_text)
        }
    }

    pub fn get_index(&self) -> usize {
        self.index.1
    }

    pub fn get_full_index(&self) -> (usize, usize) {
        self.index
    }

    pub fn get_matching_ranges(&self) -> &[(usize, usize)] {
        &self.matching_ranges
    }
}

impl Clone for Item {
    fn clone(&self) -> Item {
        Item {
            index: self.index,
            orig_text: self.orig_text.clone(),
            text: self.text.clone(),
            using_transform_fields: self.using_transform_fields,
            matching_ranges: self.matching_ranges.clone(),
            ansi_enabled: self.ansi_enabled,
        }
    }
}

pub type Rank = [i64; 4]; // score, index, start, end

#[derive(PartialEq, Eq, Clone, Debug)]
#[allow(dead_code)]
pub enum MatchedRange {
    ByteRange(usize, usize), // range of bytes
    Chars(Vec<usize>),       // individual characters matched
}

#[derive(Clone, Debug)]
pub struct MatchedItem {
    pub item: Arc<Item>,
    pub rank: Rank,
    pub matched_range: Option<MatchedRange>, // range of chars that matched the pattern
}

impl MatchedItem {
    pub fn builder(item: Arc<Item>) -> Self {
        MatchedItem {
            item,
            rank: [0, 0, 0, 0],
            matched_range: None,
        }
    }

    pub fn matched_range(mut self, range: MatchedRange) -> Self {
        self.matched_range = Some(range);
        self
    }

    pub fn rank(mut self, rank: Rank) -> Self {
        self.rank = rank;
        self
    }

    pub fn build(self) -> Self {
        self
    }
}

impl Ord for MatchedItem {
    fn cmp(&self, other: &MatchedItem) -> Ordering {
        self.rank.cmp(&other.rank)
    }
}

// `PartialOrd` needs to be implemented as well.
impl PartialOrd for MatchedItem {
    fn partial_cmp(&self, other: &MatchedItem) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for MatchedItem {
    fn eq(&self, other: &MatchedItem) -> bool {
        self.rank == other.rank
    }
}

impl Eq for MatchedItem {}

const ITEM_POOL_CAPACITY : usize = 1024;

pub struct ItemPool {
    pool: SpinLock<Vec<Arc<Item>>>,
    /// number of items that was `take`n
    taken: AtomicUsize,
}

impl ItemPool {
    pub fn new() -> Self {
        Self {
            pool: SpinLock::new(Vec::with_capacity(ITEM_POOL_CAPACITY)),
            taken: AtomicUsize::new(0),
        }
    }

    pub fn len(&self) -> usize {
        let items = self.pool.lock();
        items.len()
    }

    pub fn clear(&self) {
        let mut items = self.pool.lock();
        items.clear();
        self.taken.store(0, AtomicOrdering::SeqCst);
    }

    pub fn reset(&self) {
        // lock to ensure consistency
        let items = self.pool.lock();
        self.taken.store(0, AtomicOrdering::SeqCst);
    }

    pub fn append(&self, items: &mut Vec<Arc<Item>>) {
        let mut pool = self.pool.lock();
        pool.append(items);
    }

    pub fn take(&self) -> ItemPoolGuard<Arc<Item>> {
        let guard = self.pool.lock();
        let taken = self.taken.swap(guard.len(), AtomicOrdering::SeqCst);
        ItemPoolGuard {
            guard,
            start: taken,
        }
    }
}

pub struct ItemPoolGuard<'a, T: Sized + 'a> {
    guard: SpinLockGuard<'a, Vec<T>>,
    start: usize,
}

impl<'mutex, T: Sized> Deref for ItemPoolGuard<'mutex, T> {
    type Target = [T];

    fn deref(&self) -> &[T] {
        &self.guard[self.start..]
    }
}
<|endoftext|>


<commit_before>use crate::progress_bar::ProgressBar;
use crate::style::ProgressStyle;
use std::borrow::Cow;
use std::convert::TryFrom;
use std::io::{self, IoSliceMut};
use std::iter::FusedIterator;
use std::time::Duration;
#[cfg(feature = "tokio")]
use std::{
    pin::Pin,
    task::{Context, Poll},
};
#[cfg(feature = "tokio")]
use tokio::io::{ReadBuf, SeekFrom};

/// Wraps an iterator to display its progress.
pub trait ProgressIterator
where
    Self: Sized + Iterator,
{
    /// Wrap an iterator with default styling. Uses `Iterator::size_hint` to get length.
    /// Returns `Some(..)` only if `size_hint.1` is `Some`. If you want to create a progress bar
    /// even if `size_hint.1` returns `None` use `progress_count` or `progress_with` instead.
    fn try_progress(self) -> Option<ProgressBarIter<Self>> {
        self.size_hint()
            .1
            .map(|len| self.progress_count(u64::try_from(len).unwrap()))
    }

    /// Wrap an iterator with default styling.
    fn progress(self) -> ProgressBarIter<Self>
    where
        Self: ExactSizeIterator,
    {
        let len = u64::try_from(self.len()).unwrap();
        self.progress_count(len)
    }

    /// Wrap an iterator with an explicit element count.
    fn progress_count(self, len: u64) -> ProgressBarIter<Self> {
        self.progress_with(ProgressBar::new(len))
    }

    /// Wrap an iterator with a custom progress bar.
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>;

    /// Wrap an iterator with a progress bar and style it.
    fn progress_with_style(self, style: crate::ProgressStyle) -> ProgressBarIter<Self>
    where
        Self: ExactSizeIterator,
    {
        let len = u64::try_from(self.len()).unwrap();
        let bar = ProgressBar::new(len).with_style(style);
        self.progress_with(bar)
    }
}

/// Wraps an iterator to display its progress.
#[derive(Debug)]
pub struct ProgressBarIter<T> {
    pub(crate) it: T,
    pub progress: ProgressBar,
}

impl<T> ProgressBarIter<T> {
    /// Builder-like function for setting underlying progress bar's style.
    ///
    /// See [ProgressBar::with_style].
    pub fn with_style(mut self, style: ProgressStyle) -> ProgressBarIter<T> {
        self.progress = self.progress.with_style(style);
        self
    }

    /// Builder-like function for setting underlying progress bar's prefix.
    ///
    /// See [ProgressBar::with_prefix].
    pub fn with_prefix(mut self, prefix: impl Into<Cow<'static, str>>) -> ProgressBarIter<T> {
        self.progress = self.progress.with_prefix(prefix);
        self
    }

    /// Builder-like function for setting underlying progress bar's message.
    ///
    /// See [ProgressBar::with_message].
    pub fn with_message(mut self, message: impl Into<Cow<'static, str>>) -> ProgressBarIter<T> {
        self.progress = self.progress.with_message(message);
        self
    }

    /// Builder-like function for setting underlying progress bar's position.
    ///
    /// See [ProgressBar::with_position].
    pub fn with_position(mut self, position: u64) -> ProgressBarIter<T> {
        self.progress = self.progress.with_position(position);
        self
    }

    /// Builder-like function for setting underlying progress bar's elapsed time.
    ///
    /// See [ProgressBar::with_elapsed].
    pub fn with_elapsed(mut self, elapsed: Duration) -> ProgressBarIter<T> {
        self.progress = self.progress.with_elapsed(elapsed);
        self
    }
}

impl<S, T: Iterator<Item = S>> Iterator for ProgressBarIter<T> {
    type Item = S;

    fn next(&mut self) -> Option<Self::Item> {
        let item = self.it.next();

        if item.is_some() {
            self.progress.inc(1);
        } else if !self.progress.is_finished() {
            self.progress.finish_using_style();
        }

        item
    }
}

impl<T: ExactSizeIterator> ExactSizeIterator for ProgressBarIter<T> {
    fn len(&self) -> usize {
        self.it.len()
    }
}

impl<T: DoubleEndedIterator> DoubleEndedIterator for ProgressBarIter<T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let item = self.it.next_back();

        if item.is_some() {
            self.progress.inc(1);
        } else if !self.progress.is_finished() {
            self.progress.finish_using_style();
        }

        item
    }
}

impl<T: FusedIterator> FusedIterator for ProgressBarIter<T> {}

impl<R: io::Read> io::Read for ProgressBarIter<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let inc = self.it.read(buf)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
        let inc = self.it.read_vectored(bufs)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
        let inc = self.it.read_to_string(buf)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        self.it.read_exact(buf)?;
        self.progress.inc(buf.len() as u64);
        Ok(())
    }
}

impl<R: io::BufRead> io::BufRead for ProgressBarIter<R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.it.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        self.it.consume(amt);
        self.progress.inc(amt as u64);
    }
}

impl<S: io::Seek> io::Seek for ProgressBarIter<S> {
    fn seek(&mut self, f: io::SeekFrom) -> io::Result<u64> {
        self.it.seek(f).map(|pos| {
            self.progress.set_position(pos);
            pos
        })
    }
    // Pass this through to preserve optimizations that the inner I/O object may use here
    // Also avoid sending a set_position update when the position hasn't changed
    fn stream_position(&mut self) -> io::Result<u64> {
        self.it.stream_position()
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncWrite + Unpin> tokio::io::AsyncWrite for ProgressBarIter<W> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.it).poll_write(cx, buf).map(|poll| {
            poll.map(|inc| {
                self.progress.inc(inc as u64);
                inc
            })
        })
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.it).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.it).poll_shutdown(cx)
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncRead + Unpin> tokio::io::AsyncRead for ProgressBarIter<W> {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let prev_len = buf.filled().len() as u64;
        if let Poll::Ready(e) = Pin::new(&mut self.it).poll_read(cx, buf) {
            self.progress.inc(buf.filled().len() as u64 - prev_len);
            Poll::Ready(e)
        } else {
            Poll::Pending
        }
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncSeek + Unpin> tokio::io::AsyncSeek for ProgressBarIter<W> {
    fn start_seek(mut self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()> {
        Pin::new(&mut self.it).start_seek(position)
    }

    fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
        Pin::new(&mut self.it).poll_complete(cx)
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncBufRead + Unpin + tokio::io::AsyncRead> tokio::io::AsyncBufRead
    for ProgressBarIter<W>
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let this = self.get_mut();
        let result = Pin::new(&mut this.it).poll_fill_buf(cx);
        if let Poll::Ready(Ok(buf)) = &result {
            this.progress.inc(buf.len() as u64);
        }
        result
    }

    fn consume(mut self: Pin<&mut Self>, amt: usize) {
        Pin::new(&mut self.it).consume(amt);
    }
}

impl<W: io::Write> io::Write for ProgressBarIter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.it.write(buf).map(|inc| {
            self.progress.inc(inc as u64);
            inc
        })
    }

    fn write_vectored(&mut self, bufs: &[io::IoSlice]) -> io::Result<usize> {
        self.it.write_vectored(bufs).map(|inc| {
            self.progress.inc(inc as u64);
            inc
        })
    }

    fn flush(&mut self) -> io::Result<()> {
        self.it.flush()
    }

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.it.write_all(buf).map(|()| {
            self.progress.inc(buf.len() as u64);
        })
    }

    // write_fmt can not be captured with reasonable effort.
    // as it uses write_all internally by default that should not be a problem.
    // fn write_fmt(&mut self, fmt: fmt::Arguments) -> io::Result<()>;
}

impl<S, T: Iterator<Item = S>> ProgressIterator for T {
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self> {
        ProgressBarIter { it: self, progress }
    }
}

#[cfg(test)]
mod test {
    use crate::iter::{ProgressBarIter, ProgressIterator};
    use crate::progress_bar::ProgressBar;
    use crate::ProgressStyle;

    #[test]
    fn it_can_wrap_an_iterator() {
        let v = vec![1, 2, 3];
        let wrap = |it: ProgressBarIter<_>| {
            assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
        };

        wrap(v.iter().progress());
        wrap(v.iter().progress_count(3));
        wrap({
            let pb = ProgressBar::new(v.len() as u64);
            v.iter().progress_with(pb)
        });
        wrap({
            let style = ProgressStyle::default_bar().template("{wide_bar:.red} {percent}/100%");
            v.iter().progress_with_style(style)
        });
    }
}
<commit_msg>Normalize import style<commit_after>use std::borrow::Cow;
use std::convert::TryFrom;
use std::io::{self, IoSliceMut};
use std::iter::FusedIterator;
#[cfg(feature = "tokio")]
use std::pin::Pin;
#[cfg(feature = "tokio")]
use std::task::{Context, Poll};
use std::time::Duration;

#[cfg(feature = "tokio")]
use tokio::io::{ReadBuf, SeekFrom};

use crate::progress_bar::ProgressBar;
use crate::style::ProgressStyle;

/// Wraps an iterator to display its progress.
pub trait ProgressIterator
where
    Self: Sized + Iterator,
{
    /// Wrap an iterator with default styling. Uses `Iterator::size_hint` to get length.
    /// Returns `Some(..)` only if `size_hint.1` is `Some`. If you want to create a progress bar
    /// even if `size_hint.1` returns `None` use `progress_count` or `progress_with` instead.
    fn try_progress(self) -> Option<ProgressBarIter<Self>> {
        self.size_hint()
            .1
            .map(|len| self.progress_count(u64::try_from(len).unwrap()))
    }

    /// Wrap an iterator with default styling.
    fn progress(self) -> ProgressBarIter<Self>
    where
        Self: ExactSizeIterator,
    {
        let len = u64::try_from(self.len()).unwrap();
        self.progress_count(len)
    }

    /// Wrap an iterator with an explicit element count.
    fn progress_count(self, len: u64) -> ProgressBarIter<Self> {
        self.progress_with(ProgressBar::new(len))
    }

    /// Wrap an iterator with a custom progress bar.
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self>;

    /// Wrap an iterator with a progress bar and style it.
    fn progress_with_style(self, style: crate::ProgressStyle) -> ProgressBarIter<Self>
    where
        Self: ExactSizeIterator,
    {
        let len = u64::try_from(self.len()).unwrap();
        let bar = ProgressBar::new(len).with_style(style);
        self.progress_with(bar)
    }
}

/// Wraps an iterator to display its progress.
#[derive(Debug)]
pub struct ProgressBarIter<T> {
    pub(crate) it: T,
    pub progress: ProgressBar,
}

impl<T> ProgressBarIter<T> {
    /// Builder-like function for setting underlying progress bar's style.
    ///
    /// See [ProgressBar::with_style].
    pub fn with_style(mut self, style: ProgressStyle) -> ProgressBarIter<T> {
        self.progress = self.progress.with_style(style);
        self
    }

    /// Builder-like function for setting underlying progress bar's prefix.
    ///
    /// See [ProgressBar::with_prefix].
    pub fn with_prefix(mut self, prefix: impl Into<Cow<'static, str>>) -> ProgressBarIter<T> {
        self.progress = self.progress.with_prefix(prefix);
        self
    }

    /// Builder-like function for setting underlying progress bar's message.
    ///
    /// See [ProgressBar::with_message].
    pub fn with_message(mut self, message: impl Into<Cow<'static, str>>) -> ProgressBarIter<T> {
        self.progress = self.progress.with_message(message);
        self
    }

    /// Builder-like function for setting underlying progress bar's position.
    ///
    /// See [ProgressBar::with_position].
    pub fn with_position(mut self, position: u64) -> ProgressBarIter<T> {
        self.progress = self.progress.with_position(position);
        self
    }

    /// Builder-like function for setting underlying progress bar's elapsed time.
    ///
    /// See [ProgressBar::with_elapsed].
    pub fn with_elapsed(mut self, elapsed: Duration) -> ProgressBarIter<T> {
        self.progress = self.progress.with_elapsed(elapsed);
        self
    }
}

impl<S, T: Iterator<Item = S>> Iterator for ProgressBarIter<T> {
    type Item = S;

    fn next(&mut self) -> Option<Self::Item> {
        let item = self.it.next();

        if item.is_some() {
            self.progress.inc(1);
        } else if !self.progress.is_finished() {
            self.progress.finish_using_style();
        }

        item
    }
}

impl<T: ExactSizeIterator> ExactSizeIterator for ProgressBarIter<T> {
    fn len(&self) -> usize {
        self.it.len()
    }
}

impl<T: DoubleEndedIterator> DoubleEndedIterator for ProgressBarIter<T> {
    fn next_back(&mut self) -> Option<Self::Item> {
        let item = self.it.next_back();

        if item.is_some() {
            self.progress.inc(1);
        } else if !self.progress.is_finished() {
            self.progress.finish_using_style();
        }

        item
    }
}

impl<T: FusedIterator> FusedIterator for ProgressBarIter<T> {}

impl<R: io::Read> io::Read for ProgressBarIter<R> {
    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
        let inc = self.it.read(buf)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_vectored(&mut self, bufs: &mut [IoSliceMut<'_>]) -> io::Result<usize> {
        let inc = self.it.read_vectored(bufs)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
        let inc = self.it.read_to_string(buf)?;
        self.progress.inc(inc as u64);
        Ok(inc)
    }

    fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
        self.it.read_exact(buf)?;
        self.progress.inc(buf.len() as u64);
        Ok(())
    }
}

impl<R: io::BufRead> io::BufRead for ProgressBarIter<R> {
    fn fill_buf(&mut self) -> io::Result<&[u8]> {
        self.it.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        self.it.consume(amt);
        self.progress.inc(amt as u64);
    }
}

impl<S: io::Seek> io::Seek for ProgressBarIter<S> {
    fn seek(&mut self, f: io::SeekFrom) -> io::Result<u64> {
        self.it.seek(f).map(|pos| {
            self.progress.set_position(pos);
            pos
        })
    }
    // Pass this through to preserve optimizations that the inner I/O object may use here
    // Also avoid sending a set_position update when the position hasn't changed
    fn stream_position(&mut self) -> io::Result<u64> {
        self.it.stream_position()
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncWrite + Unpin> tokio::io::AsyncWrite for ProgressBarIter<W> {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        Pin::new(&mut self.it).poll_write(cx, buf).map(|poll| {
            poll.map(|inc| {
                self.progress.inc(inc as u64);
                inc
            })
        })
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.it).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.it).poll_shutdown(cx)
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncRead + Unpin> tokio::io::AsyncRead for ProgressBarIter<W> {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let prev_len = buf.filled().len() as u64;
        if let Poll::Ready(e) = Pin::new(&mut self.it).poll_read(cx, buf) {
            self.progress.inc(buf.filled().len() as u64 - prev_len);
            Poll::Ready(e)
        } else {
            Poll::Pending
        }
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncSeek + Unpin> tokio::io::AsyncSeek for ProgressBarIter<W> {
    fn start_seek(mut self: Pin<&mut Self>, position: SeekFrom) -> io::Result<()> {
        Pin::new(&mut self.it).start_seek(position)
    }

    fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
        Pin::new(&mut self.it).poll_complete(cx)
    }
}

#[cfg(feature = "tokio")]
impl<W: tokio::io::AsyncBufRead + Unpin + tokio::io::AsyncRead> tokio::io::AsyncBufRead
    for ProgressBarIter<W>
{
    fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
        let this = self.get_mut();
        let result = Pin::new(&mut this.it).poll_fill_buf(cx);
        if let Poll::Ready(Ok(buf)) = &result {
            this.progress.inc(buf.len() as u64);
        }
        result
    }

    fn consume(mut self: Pin<&mut Self>, amt: usize) {
        Pin::new(&mut self.it).consume(amt);
    }
}

impl<W: io::Write> io::Write for ProgressBarIter<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.it.write(buf).map(|inc| {
            self.progress.inc(inc as u64);
            inc
        })
    }

    fn write_vectored(&mut self, bufs: &[io::IoSlice]) -> io::Result<usize> {
        self.it.write_vectored(bufs).map(|inc| {
            self.progress.inc(inc as u64);
            inc
        })
    }

    fn flush(&mut self) -> io::Result<()> {
        self.it.flush()
    }

    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
        self.it.write_all(buf).map(|()| {
            self.progress.inc(buf.len() as u64);
        })
    }

    // write_fmt can not be captured with reasonable effort.
    // as it uses write_all internally by default that should not be a problem.
    // fn write_fmt(&mut self, fmt: fmt::Arguments) -> io::Result<()>;
}

impl<S, T: Iterator<Item = S>> ProgressIterator for T {
    fn progress_with(self, progress: ProgressBar) -> ProgressBarIter<Self> {
        ProgressBarIter { it: self, progress }
    }
}

#[cfg(test)]
mod test {
    use crate::iter::{ProgressBarIter, ProgressIterator};
    use crate::progress_bar::ProgressBar;
    use crate::ProgressStyle;

    #[test]
    fn it_can_wrap_an_iterator() {
        let v = vec![1, 2, 3];
        let wrap = |it: ProgressBarIter<_>| {
            assert_eq!(it.map(|x| x * 2).collect::<Vec<_>>(), vec![2, 4, 6]);
        };

        wrap(v.iter().progress());
        wrap(v.iter().progress_count(3));
        wrap({
            let pb = ProgressBar::new(v.len() as u64);
            v.iter().progress_with(pb)
        });
        wrap({
            let style = ProgressStyle::default_bar().template("{wide_bar:.red} {percent}/100%");
            v.iter().progress_with_style(style)
        });
    }
}
<|endoftext|>