use crate::Stream; use pin_project_lite::pin_project; use std::io; use std::pin::Pin; use std::task::{Context, Poll}; use tokio::io::{AsyncBufRead, Split}; pin_project! { /// A wrapper around [`tokio::io::Split`] that implements [`Stream`]. /// /// # Example /// /// ``` /// use tokio::io::AsyncBufReadExt; /// use tokio_stream::{StreamExt, wrappers::SplitStream}; /// /// # #[tokio::main(flavor = "current_thread")] /// # async fn main() -> std::io::Result<()> { /// let input = "Hello\nWorld\n".as_bytes(); /// let lines = AsyncBufReadExt::split(input, b'\n'); /// /// let mut stream = SplitStream::new(lines); /// while let Some(line) = stream.next().await { /// println!("length = {}", line?.len()) /// } /// # Ok(()) /// # } /// ``` /// [`tokio::io::Split`]: struct@tokio::io::Split /// [`Stream`]: trait@crate::Stream #[derive(Debug)] #[cfg_attr(docsrs, doc(cfg(feature = "io-util")))] pub struct SplitStream { #[pin] inner: Split, } } impl SplitStream { /// Create a new `SplitStream`. pub fn new(split: Split) -> Self { Self { inner: split } } /// Get back the inner `Split`. pub fn into_inner(self) -> Split { self.inner } /// Obtain a pinned reference to the inner `Split`. pub fn as_pin_mut(self: Pin<&mut Self>) -> Pin<&mut Split> { self.project().inner } } impl Stream for SplitStream { type Item = io::Result>; fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { self.project() .inner .poll_next_segment(cx) .map(Result::transpose) } } impl AsRef> for SplitStream { fn as_ref(&self) -> &Split { &self.inner } } impl AsMut> for SplitStream { fn as_mut(&mut self) -> &mut Split { &mut self.inner } }