initial try_stream impl

Author: carllerche

async-stream-impl/src/lib.rs

@@ -10,10 +10,10 @@ use syn::visit_mut::VisitMut;
 struct AsyncStreamImpl {
     yielder: syn::Ident,
     stmts: Vec<syn::Stmt>,
-    num_yield: u32,
 }
 
 struct Scrub {
+    is_try: bool,
     yielder: syn::Ident,
     unit: Box<syn::Expr>,
     num_yield: u32,
@@ -25,42 +25,51 @@ impl Parse for AsyncStreamImpl {
         input.parse::<Token![,]>()?;
 
         let mut stmts = vec![];
-        let mut scrub = Scrub {
-            yielder,
-            unit: syn::parse_quote!(()),
-            num_yield: 0,
-        };
 
         while !input.is_empty() {
-            let mut stmt = input.parse()?;
-            scrub.visit_stmt_mut(&mut stmt);
-            stmts.push(stmt);
+            stmts.push(input.parse()?);
         }
 
-        let Scrub { yielder, num_yield, .. } = scrub;
-
         Ok(AsyncStreamImpl {
             yielder,
             stmts,
-            num_yield,
         })
     }
 }
 
 impl VisitMut for Scrub {
     fn visit_expr_mut(&mut self, i: &mut syn::Expr) {
         match i {
-            syn::Expr::Yield(expr) => {
+            syn::Expr::Yield(yield_expr) => {
                 self.num_yield += 1;
 
-                let value_expr = if let Some(ref e) = expr.expr {
+                let value_expr = if let Some(ref e) = yield_expr.expr {
                     e
                 } else {
                     &self.unit
                 };
 
                 let ident = &self.yielder;
-                *i = syn::parse_quote! { #ident.send(#value_expr).await };
+
+                *i = if self.is_try {
+                    syn::parse_quote! { #ident.send(Ok(#value_expr)).await }
+                } else {
+                    syn::parse_quote! { #ident.send(#value_expr).await }
+                };
+            }
+            syn::Expr::Try(try_expr) => {
+                let ident = &self.yielder;
+                let e = &try_expr.expr;
+
+                *i = syn::parse_quote! {
+                    match #e {
+                        Ok(v) => v,
+                        Err(e) => {
+                            #ident.send(Err(e)).await;
+                            return;
+                        }
+                    }
+                };
             }
             expr => syn::visit_mut::visit_expr_mut(self, expr),
         }
@@ -71,11 +80,58 @@ impl VisitMut for Scrub {
 pub fn async_stream_impl(input: TokenStream) -> TokenStream {
     let AsyncStreamImpl {
         yielder,
-        stmts,
-        num_yield,
+        mut stmts,
     } = syn::parse_macro_input!(input as AsyncStreamImpl);
 
-    if num_yield == 0 {
+    let mut scrub = Scrub {
+        is_try: false,
+        yielder,
+        unit: syn::parse_quote!(()),
+        num_yield: 0,
+    };
+
+    for mut stmt in &mut stmts {
+        scrub.visit_stmt_mut(&mut stmt);
+    }
+
+    if scrub.num_yield == 0 {
+        let yielder = &scrub.yielder;
+
+        quote!({
+            if false {
+                #yielder.send(()).await;
+            }
+
+            #(#stmts)*
+        }).into()
+    } else {
+        quote!({
+            #(#stmts)*
+        }).into()
+    }
+}
+
+#[proc_macro_hack]
+pub fn async_try_stream_impl(input: TokenStream) -> TokenStream {
+    let AsyncStreamImpl {
+        yielder,
+        mut stmts,
+    } = syn::parse_macro_input!(input as AsyncStreamImpl);
+
+    let mut scrub = Scrub {
+        is_try: true,
+        yielder,
+        unit: syn::parse_quote!(()),
+        num_yield: 0,
+    };
+
+    for mut stmt in &mut stmts {
+        scrub.visit_stmt_mut(&mut stmt);
+    }
+
+    if scrub.num_yield == 0 {
+        let yielder = &scrub.yielder;
+
         quote!({
             if false {
                 #yielder.send(()).await;

async-stream/src/lib.rs

@@ -16,6 +16,10 @@ use proc_macro_hack::proc_macro_hack;
 #[proc_macro_hack]
 pub use async_stream_impl::async_stream_impl;
 
+#[doc(hidden)]
+#[proc_macro_hack]
+pub use async_stream_impl::async_try_stream_impl;
+
 /// Asynchronous stream
 #[macro_export]
 macro_rules! stream {
@@ -26,3 +30,14 @@ macro_rules! stream {
         })
     }}
 }
+
+/// Asynchronous fallible stream
+#[macro_export]
+macro_rules! try_stream {
+    ($($body:tt)*) => {{
+        let (mut __yield_tx, __yield_rx) = $crate::yielder::pair();
+        $crate::AsyncStream::new(__yield_rx, async move {
+            $crate::async_try_stream_impl!(__yield_tx, $($body)*)
+        })
+    }}
+}

async-stream/tests/basic.rs renamed to async-stream/tests/stream.rs

@@ -0,0 +1,36 @@
+#![feature(async_await)]
+
+use async_stream::try_stream;
+
+use tokio::prelude::*;
+
+#[tokio::test]
+async fn single_err() {
+    let s = try_stream! {
+        if true {
+            Err("hello")?;
+        } else {
+            yield "world";
+        }
+
+        unreachable!();
+    };
+
+    let values: Vec<_> = s.collect().await;
+    assert_eq!(1, values.len());
+    assert_eq!(Err("hello"), values[0]);
+}
+
+#[tokio::test]
+async fn yield_then_err() {
+    let s = try_stream! {
+        yield "hello";
+        Err("world")?;
+        unreachable!();
+    };
+
+    let values: Vec<_> = s.collect().await;
+    assert_eq!(2, values.len());
+    assert_eq!(Ok("hello"), values[0]);
+    assert_eq!(Err("world"), values[1]);
+}