Add async_try_stream to support ? operator in async stream

Author: taiki-e

Cargo.toml

@@ -13,6 +13,9 @@ description = """
 Async stream for Rust and the futures crate.
 """
 
+[package.metadata.docs.rs]
+all-features = true
+
 [workspace]
 members = ["futures-async-stream-macro"]
 

README.md

@@ -27,7 +27,7 @@ Add this to your `Cargo.toml`:
 ```toml
 [dependencies]
 futures-async-stream = "0.1.0-alpha.5"
-futures-preview = "0.3.0-alpha.17"
+futures-preview = "0.3.0-alpha.18"
 ```
 
 The current futures-async-stream requires Rust nightly 2019-08-21 or later.
@@ -150,10 +150,27 @@ impl Foo for Bar {
 }
 ```
 
+## `#[async_try_stream]` and `async_try_stream_block!`
+
+`?` operator can be used with the `#[async_try_stream]` and `async_try_stream_block!`. The `Item` of the returned stream is `Result` with `Ok` being the value yielded and `Err` the error type returned by `?` operator or `return Err(...)`.
+
+```rust
+#![feature(generators)]
+use futures::stream::Stream;
+use futures_async_stream::async_try_stream;
+
+#[async_try_stream(ok = i32, error = Box<dyn std::error::Error + Send + Sync>)]
+async fn foo(stream: impl Stream<Item = String>) {
+    #[for_await]
+    for x in stream {
+        yield x.parse()?;
+    }
+}
+```
+
 <!--
 ## List of features that may be added in the future as an extension of this feature:
 
-  * `async_try_stream` (https://github.com/rust-lang-nursery/futures-rs/pull/1548#discussion_r287558350)
   * `async_sink` (https://github.com/rust-lang-nursery/futures-rs/pull/1548#issuecomment-486205382)
   * Support `.await` in macro (https://github.com/rust-lang-nursery/futures-rs/pull/1548#discussion_r285341883)
   * Parallel version of `for_await` (https://github.com/rustasync/runtime/pull/25)
@@ -214,8 +231,7 @@ where
     type Item = i32;
 
     fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
-        let this = self.project();
-        if let Some(x) = ready!(this.stream.poll_next(cx)) {
+        if let Some(x) = ready!(self.project().stream.poll_next(cx)) {
             Poll::Ready(Some(x.parse().unwrap()))
         } else {
             Poll::Ready(None)

futures-async-stream-macro/Cargo.toml

@@ -12,6 +12,9 @@ description = """
 Definition of the `#[async_stream]` macro for the `futures-async-stream` crate as well as a few other assorted macros.
 """
 
+[package.metadata.docs.rs]
+all-features = true
+
 [lib]
 proc-macro = true
 

futures-async-stream-macro/src/lib.rs

@@ -22,6 +22,7 @@ mod utils;
 
 mod elision;
 mod stream;
+mod try_stream;
 mod visitor;
 
 /// Processes streams using a for loop.
@@ -40,12 +41,25 @@ pub fn for_await(args: TokenStream, input: TokenStream) -> TokenStream {
 /// Creates streams via generators.
 #[proc_macro_attribute]
 pub fn async_stream(args: TokenStream, input: TokenStream) -> TokenStream {
-    stream::async_stream(args.into(), input.into()).unwrap_or_else(|e| e.to_compile_error()).into()
+    stream::attribute(args.into(), input.into()).unwrap_or_else(|e| e.to_compile_error()).into()
 }
 
 /// Creates streams via generators.
 #[proc_macro]
 pub fn async_stream_block(input: TokenStream) -> TokenStream {
     let input = TokenStream::from(TokenTree::Group(Group::new(Delimiter::Brace, input)));
-    stream::async_stream_block(input.into()).unwrap_or_else(|e| e.to_compile_error()).into()
+    stream::block_macro(input.into()).unwrap_or_else(|e| e.to_compile_error()).into()
+}
+
+/// Creates streams via generators.
+#[proc_macro_attribute]
+pub fn async_try_stream(args: TokenStream, input: TokenStream) -> TokenStream {
+    try_stream::attribute(args.into(), input.into()).unwrap_or_else(|e| e.to_compile_error()).into()
+}
+
+/// Creates streams via generators.
+#[proc_macro]
+pub fn async_try_stream_block(input: TokenStream) -> TokenStream {
+    let input = TokenStream::from(TokenTree::Group(Group::new(Delimiter::Brace, input)));
+    try_stream::block_macro(input.into()).unwrap_or_else(|e| e.to_compile_error()).into()
 }

futures-async-stream-macro/src/stream.rs

@@ -17,7 +17,7 @@ use crate::{
 // =================================================================================================
 // async_stream
 
-pub(super) fn async_stream(args: TokenStream, input: TokenStream) -> Result<TokenStream> {
+pub(super) fn attribute(args: TokenStream, input: TokenStream) -> Result<TokenStream> {
     parse_async_stream_fn(args, input)
 }
 
@@ -38,9 +38,8 @@ impl Parse for Item {
     }
 }
 
-// TODO: rename to `ReturnType`
 #[derive(Clone, Copy)]
-enum ReturnTypeKind {
+pub(super) enum ReturnTypeKind {
     // impl Stream<Item = ..> $(+ $lifetime)?
     Default,
     // Pin<Box<dyn Stream<Item = ..> (+ Send)? $(+ $lifetime)?>>
@@ -122,16 +121,16 @@ impl Parse for Args {
     }
 }
 
-struct FnSig {
-    attrs: Vec<Attribute>,
-    vis: Visibility,
-    sig: Signature,
-    block: Block,
-    semi: Option<Token![;]>,
+pub(super) struct FnSig {
+    pub(super) attrs: Vec<Attribute>,
+    pub(super) vis: Visibility,
+    pub(super) sig: Signature,
+    pub(super) block: Block,
+    pub(super) semi: Option<Token![;]>,
 }
 
 impl FnSig {
-    fn parse(input: TokenStream, boxed: ReturnTypeKind) -> Result<Self> {
+    pub(super) fn parse(input: TokenStream, boxed: ReturnTypeKind) -> Result<Self> {
         match boxed {
             ReturnTypeKind::Default => syn::parse2(input).map(ItemFn::into),
             ReturnTypeKind::Boxed { .. } => {
@@ -167,7 +166,7 @@ impl From<TraitItemMethod> for FnSig {
     }
 }
 
-fn validate_async_stream_fn(item: &FnSig) -> Result<()> {
+pub(super) fn validate_async_stream_fn(item: &FnSig) -> Result<()> {
     if item.sig.asyncness.is_none() {
         return Err(error!(item.sig.fn_token, "async stream must be declared as async"));
     }
@@ -193,15 +192,7 @@ fn validate_async_stream_fn(item: &FnSig) -> Result<()> {
     Ok(())
 }
 
-fn parse_async_stream_fn(args: TokenStream, input: TokenStream) -> Result<TokenStream> {
-    let args: Args = syn::parse2(args)?;
-    let item = FnSig::parse(input, args.boxed)?;
-
-    validate_async_stream_fn(&item)?;
-    Ok(expand_async_stream_fn(item, &args))
-}
-
-fn expand_async_body(inputs: Punctuated<FnArg, Comma>) -> (Vec<FnArg>, Vec<Local>) {
+pub(super) fn expand_async_body(inputs: Punctuated<FnArg, Comma>) -> (Vec<FnArg>, Vec<Local>) {
     let mut arguments: Vec<FnArg> = Vec::new();
     let mut statements: Vec<Local> = Vec::new();
 
@@ -274,7 +265,13 @@ fn expand_async_body(inputs: Punctuated<FnArg, Comma>) -> (Vec<FnArg>, Vec<Local
     (arguments, statements)
 }
 
-fn make_gen_body(statements: &[Local], block: &Block, gen_function: &TokenStream) -> TokenStream {
+pub(super) fn make_gen_body(
+    statements: &[Local],
+    block: &Block,
+    gen_function: &TokenStream,
+    ret_value: &TokenStream,
+    ret_ty: &TokenStream,
+) -> TokenStream {
     let block_inner = quote! {
         #(#statements)*
         #block
@@ -292,7 +289,7 @@ fn make_gen_body(statements: &[Local], block: &Block, gen_function: &TokenStream
         // have any `yield` statements.
         #[allow(unreachable_code)]
         {
-            return;
+            return #ret_value;
             loop { yield ::futures_async_stream::reexport::task::Poll::Pending }
         }
     };
@@ -302,10 +299,18 @@ fn make_gen_body(statements: &[Local], block: &Block, gen_function: &TokenStream
     });
 
     quote! {
-        #gen_function (static move || -> () #gen_body)
+        #gen_function(static move || -> #ret_ty #gen_body)
     }
 }
 
+fn parse_async_stream_fn(args: TokenStream, input: TokenStream) -> Result<TokenStream> {
+    let args: Args = syn::parse2(args)?;
+    let item = FnSig::parse(input, args.boxed)?;
+
+    validate_async_stream_fn(&item)?;
+    Ok(expand_async_stream_fn(item, &args))
+}
+
 fn expand_async_stream_fn(item: FnSig, args: &Args) -> TokenStream {
     let FnSig { attrs, vis, sig, mut block, semi } = item;
     let Signature { unsafety, abi, fn_token, ident, mut generics, inputs, .. } = sig;
@@ -324,7 +329,7 @@ fn expand_async_stream_fn(item: FnSig, args: &Args) -> TokenStream {
     let output_span = first_last(item);
     let gen_function = quote!(::futures_async_stream::stream::from_generator);
     let gen_function = respan(gen_function, output_span);
-    let mut body_inner = make_gen_body(&statements, &block, &gen_function);
+    let mut body_inner = make_gen_body(&statements, &block, &gen_function, &quote!(), &quote!(()));
 
     if let ReturnTypeKind::Boxed { .. } = args.boxed {
         let body = quote! { ::futures_async_stream::reexport::boxed::Box::pin(#body_inner) };
@@ -344,15 +349,15 @@ fn expand_async_stream_fn(item: FnSig, args: &Args) -> TokenStream {
             // Raw `impl` breaks syntax highlighting in some editors.
             let impl_token = token::Impl::default();
             quote! {
-                #impl_token ::futures_async_stream::stream::Stream<Item = #item> + #(#lifetimes +)*
+                #impl_token ::futures_async_stream::reexport::Stream<Item = #item> + #(#lifetimes +)*
             }
         }
         ReturnTypeKind::Boxed { send } => {
             let send = if send { Some(quote!(+ Send)) } else { None };
             quote! {
                 ::futures_async_stream::reexport::pin::Pin<
                     ::futures_async_stream::reexport::boxed::Box<
-                        dyn ::futures_async_stream::stream::Stream<Item = #item> #send + #(#lifetimes +)*
+                        dyn ::futures_async_stream::reexport::Stream<Item = #item> #send + #(#lifetimes +)*
                     >
                 >
             }
@@ -375,13 +380,13 @@ fn expand_async_stream_fn(item: FnSig, args: &Args) -> TokenStream {
 // =================================================================================================
 // async_stream_block
 
-pub(super) fn async_stream_block(input: TokenStream) -> Result<TokenStream> {
+pub(super) fn block_macro(input: TokenStream) -> Result<TokenStream> {
     syn::parse2(input).map(expand_async_stream_block)
 }
 
 fn expand_async_stream_block(mut expr: Expr) -> TokenStream {
     Visitor::new(Stream).visit_expr_mut(&mut expr);
 
     let gen_function = quote!(::futures_async_stream::stream::from_generator);
-    make_gen_body(&[], &block(vec![Stmt::Expr(expr)]), &gen_function)
+    make_gen_body(&[], &block(vec![Stmt::Expr(expr)]), &gen_function, &quote!(), &quote!(()))
 }