Skip to content

Commit e46ea52

Browse files
authored
Implemented TryStreamExt::try_ready_chunks (#2757)
1 parent bdd0f82 commit e46ea52

File tree

3 files changed

+182
-1
lines changed

3 files changed

+182
-1
lines changed

futures-util/src/stream/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ pub use self::try_stream::{TryBufferUnordered, TryBuffered, TryFlattenUnordered}
7171
pub use self::try_stream::TryForward;
7272

7373
#[cfg(feature = "alloc")]
74-
pub use self::try_stream::{TryChunks, TryChunksError};
74+
pub use self::try_stream::{TryChunks, TryChunksError, TryReadyChunks, TryReadyChunksError};
7575

7676
// Primitive streams
7777

futures-util/src/stream/try_stream/mod.rs

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,12 @@ mod try_chunks;
121121
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
122122
pub use self::try_chunks::{TryChunks, TryChunksError};
123123

124+
#[cfg(feature = "alloc")]
125+
mod try_ready_chunks;
126+
#[cfg(feature = "alloc")]
127+
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
128+
pub use self::try_ready_chunks::{TryReadyChunks, TryReadyChunksError};
129+
124130
mod try_unfold;
125131
#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
126132
pub use self::try_unfold::{try_unfold, TryUnfold};
@@ -557,6 +563,55 @@ pub trait TryStreamExt: TryStream {
557563
)
558564
}
559565

566+
/// An adaptor for chunking up successful, ready items of the stream inside a vector.
567+
///
568+
/// This combinator will attempt to pull successful items from this stream and buffer
569+
/// them into a local vector. At most `capacity` items will get buffered
570+
/// before they're yielded from the returned stream. If the underlying stream
571+
/// returns `Poll::Pending`, and the collected chunk is not empty, it will
572+
/// be immidiatly returned.
573+
///
574+
/// Note that the vectors returned from this iterator may not always have
575+
/// `capacity` elements. If the underlying stream ended and only a partial
576+
/// vector was created, it'll be returned. Additionally if an error happens
577+
/// from the underlying stream then the currently buffered items will be
578+
/// yielded.
579+
///
580+
/// This method is only available when the `std` or `alloc` feature of this
581+
/// library is activated, and it is activated by default.
582+
///
583+
/// This function is similar to
584+
/// [`StreamExt::ready_chunks`](crate::stream::StreamExt::ready_chunks) but exits
585+
/// early if an error occurs.
586+
///
587+
/// # Examples
588+
///
589+
/// ```
590+
/// # futures::executor::block_on(async {
591+
/// use futures::stream::{self, TryReadyChunksError, TryStreamExt};
592+
///
593+
/// let stream = stream::iter(vec![Ok::<i32, i32>(1), Ok(2), Ok(3), Err(4), Ok(5), Ok(6)]);
594+
/// let mut stream = stream.try_ready_chunks(2);
595+
///
596+
/// assert_eq!(stream.try_next().await, Ok(Some(vec![1, 2])));
597+
/// assert_eq!(stream.try_next().await, Err(TryReadyChunksError(vec![3], 4)));
598+
/// assert_eq!(stream.try_next().await, Ok(Some(vec![5, 6])));
599+
/// # })
600+
/// ```
601+
///
602+
/// # Panics
603+
///
604+
/// This method will panic if `capacity` is zero.
605+
#[cfg(feature = "alloc")]
606+
fn try_ready_chunks(self, capacity: usize) -> TryReadyChunks<Self>
607+
where
608+
Self: Sized,
609+
{
610+
assert_stream::<Result<Vec<Self::Ok>, TryReadyChunksError<Self::Ok, Self::Error>>, _>(
611+
TryReadyChunks::new(self, capacity),
612+
)
613+
}
614+
560615
/// Attempt to filter the values produced by this stream according to the
561616
/// provided asynchronous closure.
562617
///
Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
use crate::stream::{Fuse, IntoStream, StreamExt};
2+
3+
use alloc::vec::Vec;
4+
use core::fmt;
5+
use core::pin::Pin;
6+
use futures_core::stream::{FusedStream, Stream, TryStream};
7+
use futures_core::task::{Context, Poll};
8+
#[cfg(feature = "sink")]
9+
use futures_sink::Sink;
10+
use pin_project_lite::pin_project;
11+
12+
pin_project! {
13+
/// Stream for the [`try_ready_chunks`](super::TryStreamExt::try_ready_chunks) method.
14+
#[derive(Debug)]
15+
#[must_use = "streams do nothing unless polled"]
16+
pub struct TryReadyChunks<St: TryStream> {
17+
#[pin]
18+
stream: Fuse<IntoStream<St>>,
19+
cap: usize, // https://github.com/rust-lang/futures-rs/issues/1475
20+
}
21+
}
22+
23+
impl<St: TryStream> TryReadyChunks<St> {
24+
pub(super) fn new(stream: St, capacity: usize) -> Self {
25+
assert!(capacity > 0);
26+
27+
Self { stream: IntoStream::new(stream).fuse(), cap: capacity }
28+
}
29+
30+
delegate_access_inner!(stream, St, (. .));
31+
}
32+
33+
type TryReadyChunksStreamError<St> =
34+
TryReadyChunksError<<St as TryStream>::Ok, <St as TryStream>::Error>;
35+
36+
impl<St: TryStream> Stream for TryReadyChunks<St> {
37+
type Item = Result<Vec<St::Ok>, TryReadyChunksStreamError<St>>;
38+
39+
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
40+
let mut this = self.as_mut().project();
41+
42+
let mut items: Vec<St::Ok> = Vec::new();
43+
44+
loop {
45+
match this.stream.as_mut().poll_next(cx) {
46+
// Flush all the collected data if the underlying stream doesn't
47+
// contain more ready values
48+
Poll::Pending => {
49+
return if items.is_empty() {
50+
Poll::Pending
51+
} else {
52+
Poll::Ready(Some(Ok(items)))
53+
}
54+
}
55+
56+
// Push the ready item into the buffer and check whether it is full.
57+
// If so, return the buffer.
58+
Poll::Ready(Some(Ok(item))) => {
59+
if items.is_empty() {
60+
items.reserve_exact(*this.cap);
61+
}
62+
items.push(item);
63+
if items.len() >= *this.cap {
64+
return Poll::Ready(Some(Ok(items)));
65+
}
66+
}
67+
68+
// Return the already collected items and the error.
69+
Poll::Ready(Some(Err(e))) => {
70+
return Poll::Ready(Some(Err(TryReadyChunksError(items, e))));
71+
}
72+
73+
// Since the underlying stream ran out of values, return what we
74+
// have buffered, if we have anything.
75+
Poll::Ready(None) => {
76+
let last = if items.is_empty() { None } else { Some(Ok(items)) };
77+
return Poll::Ready(last);
78+
}
79+
}
80+
}
81+
}
82+
83+
fn size_hint(&self) -> (usize, Option<usize>) {
84+
let (lower, upper) = self.stream.size_hint();
85+
let lower = lower / self.cap;
86+
(lower, upper)
87+
}
88+
}
89+
90+
impl<St: TryStream + FusedStream> FusedStream for TryReadyChunks<St> {
91+
fn is_terminated(&self) -> bool {
92+
self.stream.is_terminated()
93+
}
94+
}
95+
96+
// Forwarding impl of Sink from the underlying stream
97+
#[cfg(feature = "sink")]
98+
impl<S, Item> Sink<Item> for TryReadyChunks<S>
99+
where
100+
S: TryStream + Sink<Item>,
101+
{
102+
type Error = <S as Sink<Item>>::Error;
103+
104+
delegate_sink!(stream, Item);
105+
}
106+
107+
/// Error indicating, that while chunk was collected inner stream produced an error.
108+
///
109+
/// Contains all items that were collected before an error occurred, and the stream error itself.
110+
#[derive(PartialEq, Eq)]
111+
pub struct TryReadyChunksError<T, E>(pub Vec<T>, pub E);
112+
113+
impl<T, E: fmt::Debug> fmt::Debug for TryReadyChunksError<T, E> {
114+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
115+
self.1.fmt(f)
116+
}
117+
}
118+
119+
impl<T, E: fmt::Display> fmt::Display for TryReadyChunksError<T, E> {
120+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121+
self.1.fmt(f)
122+
}
123+
}
124+
125+
#[cfg(feature = "std")]
126+
impl<T, E: fmt::Debug + fmt::Display> std::error::Error for TryReadyChunksError<T, E> {}

0 commit comments

Comments
 (0)