Merge pull request #57 from andreytkachenko/await-item-macro

`await_item` macro for streams

futures-await-await-macro/src/lib.rs

@@ -27,6 +27,30 @@ macro_rules! await {
     })
 }
 
+///
+/// Await an item from the stream
+/// Basically it does same as `await` macro, but for streams
+///
+
+#[macro_export]
+macro_rules! await_item {
+    ($e:expr) => ({
+        loop {
+            match ::futures::Stream::poll(&mut $e) {
+                ::futures::__rt::std::result::Result::Ok(::futures::Async::Ready(e)) => {
+                    break ::futures::__rt::std::result::Result::Ok(e)
+                }
+                ::futures::__rt::std::result::Result::Ok(::futures::Async::NotReady) => {}
+                ::futures::__rt::std::result::Result::Err(e) => {
+                    break ::futures::__rt::std::result::Result::Err(e)
+                }
+            }
+
+            yield ::futures::Async::NotReady
+        }
+    })
+}
+
 // TODO: This macro needs to use an extra temporary variable because of
 // rust-lang/rust#44197, once that's fixed this should just use $e directly
 // inside the yield expression

src/lib.rs

@@ -26,7 +26,7 @@ pub use futures::*;
 pub mod prelude {
     pub use futures::prelude::*;
     pub use async_macro::{async, async_stream, async_block, async_stream_block};
-    pub use await_macro::{await, stream_yield};
+    pub use await_macro::{await, stream_yield, await_item};
 }
 
 /// A hidden module that's the "runtime support" for the async/await syntax.

tests/smoke.rs

@@ -174,6 +174,24 @@ pub fn _stream8() -> Result<(), i32> {
 //     }
 // }
 
+#[async_stream(item = u64)]
+fn await_item_stream() -> Result<(), i32> {
+    stream_yield!(0);
+    stream_yield!(1);
+    Ok(())
+}
+
+#[async]
+fn test_await_item() -> Result<(), ()> {
+    let mut stream = await_item_stream();
+
+    assert_eq!(await_item!(stream), Ok(Some(0)));
+    assert_eq!(await_item!(stream), Ok(Some(1)));
+    assert_eq!(await_item!(stream), Ok(None));
+
+    Ok(())
+}
+
 #[test]
 fn main() {
     assert_eq!(foo().wait(), Ok(1));
@@ -184,6 +202,7 @@ fn main() {
     assert_eq!(_foo6(8).wait(), Err(8));
     // assert_eq!(A(11).a_foo().wait(), Ok(11));
     assert_eq!(loop_in_loop().wait(), Ok(true));
+    assert_eq!(test_await_item().wait(), Ok(()));
 }
 
 #[async]