Merge pull request #545 from ishitatsuyuki/stream_pollfn

Implement stream::poll_fn

Author: alexcrichton

src/stream/mod.rs

@@ -43,6 +43,7 @@ mod merge;
 mod once;
 mod or_else;
 mod peek;
+mod poll_fn;
 mod select;
 mod skip;
 mod skip_while;
@@ -70,6 +71,7 @@ pub use self::merge::{Merge, MergedItem};
 pub use self::once::{Once, once};
 pub use self::or_else::OrElse;
 pub use self::peek::Peekable;
+pub use self::poll_fn::{poll_fn, PollFn};
 pub use self::select::Select;
 pub use self::skip::Skip;
 pub use self::skip_while::SkipWhile;

src/stream/poll_fn.rs

@@ -0,0 +1,49 @@
+//! Definition of the `PollFn` combinator
+
+use {Stream, Poll};
+
+/// A stream which adapts a function returning `Poll`.
+///
+/// Created by the `poll_fn` function.
+#[derive(Debug)]
+#[must_use = "streams do nothing unless polled"]
+pub struct PollFn<F> {
+    inner: F,
+}
+
+/// Creates a new stream wrapping around a function returning `Poll`.
+///
+/// Polling the returned stream delegates to the wrapped function.
+///
+/// # Examples
+///
+/// ```
+/// use futures::stream::poll_fn;
+/// use futures::{Async, Poll};
+///
+/// let mut counter = 1usize;
+///
+/// let read_stream = poll_fn(move || -> Poll<Option<String>, std::io::Error> {
+///     if counter == 0 { return Ok(Async::Ready(None)); }
+///     counter -= 1;
+///     Ok(Async::Ready(Some("Hello, World!".to_owned())))
+/// });
+/// ```
+pub fn poll_fn<T, E, F>(f: F) -> PollFn<F>
+where
+    F: FnMut() -> Poll<Option<T>, E>,
+{
+    PollFn { inner: f }
+}
+
+impl<T, E, F> Stream for PollFn<F>
+where
+    F: FnMut() -> Poll<Option<T>, E>,
+{
+    type Item = T;
+    type Error = E;
+
+    fn poll(&mut self) -> Poll<Option<T>, E> {
+        (self.inner)()
+    }
+}

tests/stream.rs

@@ -1,10 +1,10 @@
 #[macro_use]
 extern crate futures;
 
-use futures::{Poll, Future, Stream, Sink};
+use futures::{Async, Future, Poll, Sink, Stream};
 use futures::executor;
-use futures::future::{ok, err};
-use futures::stream::{empty, iter, Peekable, BoxStream};
+use futures::future::{err, ok};
+use futures::stream::{empty, iter, poll_fn, BoxStream, Peekable};
 use futures::sync::oneshot;
 use futures::sync::mpsc;
 
@@ -356,3 +356,18 @@ fn concat2() {
     let c = empty::<Vec<()>, ()>();
     assert_done(move || c.concat2(), Ok(vec![]))
 }
+
+#[test]
+fn stream_poll_fn() {
+    let mut counter = 5usize;
+
+    let read_stream = poll_fn(move || -> Poll<Option<usize>, std::io::Error> {
+        if counter == 0 {
+            return Ok(Async::Ready(None));
+        }
+        counter -= 1;
+        Ok(Async::Ready(Some(counter)))
+    });
+
+    assert_eq!(read_stream.wait().count(), 5);
+}