refactor: Body::wrap_stream -> Body::from_stream

omjadas · Apr 29, 2024 · 1279433 · 1279433
1 parent de4bebe
commit 1279433
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 18 deletions.
diff --git a/benches/decoder.rs b/benches/decoder.rs
@@ -19,13 +19,13 @@ fn raw_body() -> Body {
 
 fn gzip_body() -> Body {
     let encoder = GzipEncoder::new(&BODY[..]);
-    Body::wrap_stream(ReaderStream::new(encoder))
+    Body::from_stream(ReaderStream::new(encoder))
 }
 
 fn gzip_brotli_body() -> Body {
     let encoder = GzipEncoder::new(&BODY[..]);
     let encoder = BrotliEncoder::new(BufReader::new(encoder));
-    Body::wrap_stream(ReaderStream::new(encoder))
+    Body::from_stream(ReaderStream::new(encoder))
 }
 
 fn raw_request() -> Request<Body> {

diff --git a/src/body.rs b/src/body.rs
@@ -1,12 +1,12 @@
 use crate::Error;
-use futures::{Stream, StreamExt};
+use futures::{Stream, TryStream, TryStreamExt};
 use http_body_util::{combinators::BoxBody, Collected, Empty, Full, StreamBody};
 use hyper::body::{Body as HttpBody, Bytes, Frame, Incoming, SizeHint};
 use std::pin::Pin;
 
 #[derive(Debug)]
 enum Internal {
-    BoxBody(BoxBody<Bytes, crate::Error>),
+    BoxBody(BoxBody<Bytes, Error>),
     Collected(Collected<Bytes>),
     Empty(Empty<Bytes>),
     Full(Full<Bytes>),
@@ -21,23 +21,26 @@ pub struct Body {
 }
 
 impl Body {
-    pub fn wrap_stream<S, O, E>(stream: S) -> Self
+    pub fn from_stream<S>(stream: S) -> Self
     where
-        S: Stream<Item = Result<O, E>> + Send + Sync + 'static,
-        O: Into<Bytes>,
-        E: Into<Error>,
+        S: TryStream + Send + Sync + 'static,
+        S::Ok: Into<Bytes>,
+        S::Error: Into<Error>,
     {
         Self {
             inner: Internal::BoxBody(BoxBody::new(StreamBody::new(
-                stream.map(|res| res.map(Into::into).map(Frame::data).map_err(Into::into)),
+                stream
+                    .map_ok(Into::into)
+                    .map_ok(Frame::data)
+                    .map_err(Into::into),
             ))),
         }
     }
 }
 
 impl HttpBody for Body {
     type Data = Bytes;
-    type Error = crate::Error;
+    type Error = Error;
 
     fn poll_frame(
         mut self: std::pin::Pin<&mut Self>,
@@ -76,8 +79,8 @@ impl HttpBody for Body {
     }
 }
 
-impl From<BoxBody<Bytes, crate::Error>> for Body {
-    fn from(value: BoxBody<Bytes, crate::Error>) -> Self {
+impl From<BoxBody<Bytes, Error>> for Body {
+    fn from(value: BoxBody<Bytes, Error>) -> Self {
         Self {
             inner: Internal::BoxBody(value),
         }

diff --git a/src/decoder.rs b/src/decoder.rs
@@ -70,7 +70,7 @@ impl From<Decoder<Body>> for Body {
     fn from(decoder: Decoder<Body>) -> Body {
         match decoder {
             Decoder::Body(body) => body,
-            Decoder::Decoder(decoder) => Body::wrap_stream(ReaderStream::new(decoder)),
+            Decoder::Decoder(decoder) => Body::from_stream(ReaderStream::new(decoder)),
         }
     }
 }
@@ -307,7 +307,7 @@ mod tests {
         async fn single_encoding() {
             let content = b"hello, world";
             let encoder = GzipEncoder::new(&content[..]);
-            let body = Body::wrap_stream(ReaderStream::new(encoder));
+            let body = Body::from_stream(ReaderStream::new(encoder));
 
             assert_eq!(
                 &to_bytes(decode_body(vec![&b"gzip"[..]], body).unwrap()).await[..],
@@ -320,7 +320,7 @@ mod tests {
             let content = b"hello, world";
             let encoder = GzipEncoder::new(&content[..]);
             let encoder = BrotliEncoder::new(BufReader::new(encoder));
-            let body = Body::wrap_stream(ReaderStream::new(encoder));
+            let body = Body::from_stream(ReaderStream::new(encoder));
 
             assert_eq!(
                 &to_bytes(decode_body(vec![&b"br"[..], &b"gzip"[..]], body).unwrap()).await[..],
@@ -347,7 +347,7 @@ mod tests {
             let req = Request::builder()
                 .header(CONTENT_LENGTH, 123)
                 .header(CONTENT_ENCODING, "gzip")
-                .body(Body::wrap_stream(ReaderStream::new(encoder)))
+                .body(Body::from_stream(ReaderStream::new(encoder)))
                 .unwrap();
 
             let req = decode_request(req).unwrap();
@@ -369,7 +369,7 @@ mod tests {
             let res = Response::builder()
                 .header(CONTENT_LENGTH, 123)
                 .header(CONTENT_ENCODING, "gzip")
-                .body(Body::wrap_stream(ReaderStream::new(encoder)))
+                .body(Body::from_stream(ReaderStream::new(encoder)))
                 .unwrap();
 
             let res = decode_response(res).unwrap();

diff --git a/tests/common/mod.rs b/tests/common/mod.rs
@@ -62,7 +62,7 @@ async fn test_server(req: Request<Incoming>) -> Result<Response<Body>, Infallibl
         (&Method::GET, "/hello/gzip") => Ok(Response::builder()
             .header(CONTENT_ENCODING, "gzip")
             .status(StatusCode::OK)
-            .body(Body::wrap_stream(ReaderStream::new(GzipEncoder::new(
+            .body(Body::from_stream(ReaderStream::new(GzipEncoder::new(
                 HELLO_WORLD.as_bytes(),
             ))))
             .unwrap()),