actix_web_lab/
display_stream.rs

1use std::{error::Error as StdError, fmt, io::Write as _};
2
3use actix_web::{
4    HttpResponse, Responder,
5    body::{BodyStream, MessageBody},
6};
7use bytes::{Bytes, BytesMut};
8use futures_util::{Stream, TryStreamExt as _};
9use pin_project_lite::pin_project;
10
11use crate::util::{InfallibleStream, MutWriter};
12
13pin_project! {
14    /// A buffered line formatting body stream.
15    ///
16    /// Each item yielded by the stream will be written to the response body using its
17    /// `Display` implementation.
18    ///
19    /// This has significant memory efficiency advantages over returning an array of lines when the
20    /// data set is very large because it avoids buffering the entire response.
21    ///
22    /// # Examples
23    /// ```
24    /// # use actix_web::Responder;
25    /// # use actix_web_lab::respond::DisplayStream;
26    /// # use futures_util::Stream;
27    /// fn streaming_data_source() -> impl Stream<Item = u32> {
28    ///     // get item stream from source
29    ///     # futures_util::stream::empty()
30    /// }
31    ///
32    /// async fn handler() -> impl Responder {
33    ///     let data_stream = streaming_data_source();
34    ///
35    ///     DisplayStream::new_infallible(data_stream)
36    ///         .into_responder()
37    /// }
38    /// ```
39    pub struct DisplayStream<S> {
40        // The wrapped item stream.
41        #[pin]
42        stream: S,
43    }
44}
45
46impl<S> DisplayStream<S> {
47    /// Constructs a new `DisplayStream` from a stream of lines.
48    pub fn new(stream: S) -> Self {
49        Self { stream }
50    }
51}
52
53impl<S> DisplayStream<S> {
54    /// Constructs a new `DisplayStream` from an infallible stream of lines.
55    pub fn new_infallible(stream: S) -> DisplayStream<InfallibleStream<S>> {
56        DisplayStream::new(InfallibleStream::new(stream))
57    }
58}
59
60impl<S, T, E> DisplayStream<S>
61where
62    S: Stream<Item = Result<T, E>>,
63    T: fmt::Display,
64    E: Into<Box<dyn StdError>> + 'static,
65{
66    /// Creates a chunked body stream that serializes as CSV on-the-fly.
67    pub fn into_body_stream(self) -> impl MessageBody {
68        BodyStream::new(self.into_chunk_stream())
69    }
70
71    /// Creates a `Responder` type with a line-by-line serializing stream and `text/plain`
72    /// content-type header.
73    pub fn into_responder(self) -> impl Responder
74    where
75        S: 'static,
76        T: 'static,
77        E: 'static,
78    {
79        HttpResponse::Ok()
80            .content_type(mime::TEXT_PLAIN_UTF_8)
81            .message_body(self.into_body_stream())
82            .unwrap()
83    }
84
85    /// Creates a stream of serialized chunks.
86    pub fn into_chunk_stream(self) -> impl Stream<Item = Result<Bytes, E>> {
87        self.stream.map_ok(write_display)
88    }
89}
90
91fn write_display(item: impl fmt::Display) -> Bytes {
92    let mut buf = BytesMut::new();
93    let mut wrt = MutWriter(&mut buf);
94
95    writeln!(wrt, "{item}").unwrap();
96
97    buf.freeze()
98}
99
100#[cfg(test)]
101mod tests {
102    use std::error::Error as StdError;
103
104    use actix_web::body;
105    use futures_util::stream;
106
107    use super::*;
108
109    #[actix_web::test]
110    async fn serializes_into_body() {
111        let ndjson_body = DisplayStream::new_infallible(stream::iter([123, 789, 345, 901, 456]))
112            .into_body_stream();
113
114        let body_bytes = body::to_bytes(ndjson_body)
115            .await
116            .map_err(Into::<Box<dyn StdError>>::into)
117            .unwrap();
118
119        const EXP_BYTES: &str = "123\n\
120        789\n\
121        345\n\
122        901\n\
123        456\n";
124
125        assert_eq!(body_bytes, EXP_BYTES);
126    }
127}
actix_web_lab/display_stream.rs

actix_web_lab/
display_stream.rs
