use futures01_crate::{Async, Stream};
use std::error;
use std::marker::PhantomData;
use {ErrorCompat, IntoError};
pub trait StreamExt: Stream + Sized {
fn context<C, E>(self, context: C) -> Context<Self, C, E>
where
C: IntoError<E, Source = Self::Error> + Clone,
E: error::Error + ErrorCompat;
fn with_context<F, C, E>(self, context: F) -> WithContext<Self, F, E>
where
F: FnMut() -> C,
C: IntoError<E, Source = Self::Error>,
E: error::Error + ErrorCompat;
}
impl<St> StreamExt for St
where
St: Stream,
{
fn context<C, E>(self, context: C) -> Context<Self, C, E>
where
C: IntoError<E, Source = Self::Error> + Clone,
E: error::Error + ErrorCompat,
{
Context {
stream: self,
context,
_e: PhantomData,
}
}
fn with_context<F, C, E>(self, context: F) -> WithContext<Self, F, E>
where
F: FnMut() -> C,
C: IntoError<E, Source = Self::Error>,
E: error::Error + ErrorCompat,
{
WithContext {
stream: self,
context,
_e: PhantomData,
}
}
}
pub struct Context<St, C, E> {
stream: St,
context: C,
_e: PhantomData<E>,
}
impl<St, C, E> Stream for Context<St, C, E>
where
St: Stream,
C: IntoError<E, Source = St::Error> + Clone,
E: error::Error + ErrorCompat,
{
type Item = St::Item;
type Error = E;
fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
self.stream
.poll()
.map_err(|error| self.context.clone().into_error(error))
}
}
pub struct WithContext<St, F, E> {
stream: St,
context: F,
_e: PhantomData<E>,
}
impl<St, F, C, E> Stream for WithContext<St, F, E>
where
St: Stream,
F: FnMut() -> C,
C: IntoError<E, Source = St::Error>,
E: error::Error + ErrorCompat,
{
type Item = St::Item;
type Error = E;
fn poll(&mut self) -> Result<Async<Option<Self::Item>>, Self::Error> {
self.stream.poll().map_err(|error| {
let context = &mut self.context;
context().into_error(error)
})
}
}