use std::pin::Pin;
use std::task::{Context, Poll};
use bytes::Bytes;
use futures_core::Stream;
use http::{HeaderMap, Response, StatusCode};
use http_body::{Body, Frame, SizeHint};
use super::RenderError;
pub type RenderedTemplateStream =
Pin<Box<dyn Stream<Item = Result<Frame<Bytes>, RenderError>> + Send + 'static>>;
pub enum RenderedTemplateBody {
Full(Bytes),
Stream(RenderedTemplateStream),
}
impl Body for RenderedTemplateBody {
type Data = Bytes;
type Error = RenderError;
fn poll_frame(
self: Pin<&mut Self>,
context: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
match self.get_mut() {
Self::Full(bytes) if bytes.is_empty() => Poll::Ready(None),
Self::Full(bytes) => Poll::Ready(Some(Ok(Frame::data(std::mem::take(bytes))))),
Self::Stream(stream) => stream.as_mut().poll_next(context),
}
}
fn is_end_stream(&self) -> bool {
matches!(self, Self::Full(bytes) if bytes.is_empty())
}
fn size_hint(&self) -> SizeHint {
match self {
Self::Full(bytes) => SizeHint::with_exact(bytes.len() as u64),
Self::Stream(_) => SizeHint::default(),
}
}
}
pub struct RenderedTemplate {
status: StatusCode,
headers: HeaderMap,
body: RenderedTemplateBody,
}
impl RenderedTemplate {
#[must_use]
pub fn new(headers: HeaderMap, body: RenderedTemplateBody) -> Self {
Self {
status: StatusCode::OK,
headers,
body,
}
}
#[must_use]
pub const fn get_status(&self) -> StatusCode {
self.status
}
#[must_use]
pub fn with_status(mut self, status: StatusCode) -> Self {
self.status = status;
self
}
#[must_use]
pub const fn get_headers(&self) -> &HeaderMap {
&self.headers
}
#[must_use]
pub fn get_headers_mut(&mut self) -> &mut HeaderMap {
&mut self.headers
}
#[must_use]
pub const fn get_body(&self) -> &RenderedTemplateBody {
&self.body
}
#[must_use]
pub fn into_parts(self) -> (StatusCode, HeaderMap, RenderedTemplateBody) {
(self.status, self.headers, self.body)
}
#[must_use]
pub fn into_http_response(self) -> Response<RenderedTemplateBody> {
let mut response = Response::new(self.body);
*response.status_mut() = self.status;
*response.headers_mut() = self.headers;
response
}
}