use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};
use connectrpc::client::{ClientTransport, HttpClient, ServerStream};
use futures_util::{TryFutureExt as _, TryStreamExt as _};
use crate::Error;
type Body = <HttpClient as ClientTransport>::ResponseBody;
#[must_use = "a stream makes its call when it is first read"]
pub struct ResponseStream<T> {
responses: Pin<Box<dyn futures_core::Stream<Item = Result<T, Error>> + Send>>,
}
impl<T> ResponseStream<T> {
pub async fn next(&mut self) -> Option<Result<T, Error>> {
std::future::poll_fn(|context| self.responses.as_mut().poll_next(context)).await
}
pub(crate) fn new<Call, View, Wire>(call: Call) -> Self
where
Call: Future<Output = Result<ServerStream<Body, View>, connectrpc::ConnectError>> + Send + 'static,
View: buffa::MessageView<'static, Owned = Wire> + Send + 'static,
Wire: buffa::HasMessageView<View<'static> = View> + buffa::Message + connectrpc::codec::JsonDeserialize,
T: From<Wire> + 'static,
{
let responses = call
.map_ok(|stream| {
futures_util::stream::try_unfold(stream, |mut stream| async move {
while let Some(message) = stream.message().await? {
if !message.bytes().is_empty() {
return Ok(Some((T::from(message.to_owned_message()), stream)));
}
}
Ok(None)
})
})
.try_flatten_stream()
.map_err(Error::from_wire);
Self {
responses: Box::pin(responses),
}
}
}
impl<T> futures_core::Stream for ResponseStream<T> {
type Item = Result<T, Error>;
fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.responses.as_mut().poll_next(context)
}
}
impl<T> fmt::Debug for ResponseStream<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ResponseStream").finish_non_exhaustive()
}
}