Skip to content

Commit c1d4aef

Browse files
committed
for await
1 parent 83fdd7e commit c1d4aef

File tree

2 files changed

+43
-8
lines changed

2 files changed

+43
-8
lines changed

async-stream-impl/src/lib.rs

Lines changed: 42 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
extern crate proc_macro;
22

33
use proc_macro::{TokenStream, TokenTree};
4-
use proc_macro2::Span;
4+
use proc_macro2::{Group, Span, TokenStream as TokenStream2, TokenTree as TokenTree2};
55
use quote::quote;
66
use syn::visit_mut::VisitMut;
77

@@ -19,7 +19,7 @@ struct AsyncStreamEnumHack {
1919
}
2020

2121
impl AsyncStreamEnumHack {
22-
fn parse(input: TokenStream) -> Self {
22+
fn parse(input: TokenStream) -> syn::Result<Self> {
2323
macro_rules! n {
2424
($i:ident) => {
2525
$i.next().unwrap()
@@ -44,14 +44,15 @@ impl AsyncStreamEnumHack {
4444
n!(braces); // !
4545

4646
let inner = n!(braces);
47-
let syn::Block { stmts, .. } = syn::parse(inner.clone().into()).unwrap();
47+
let inner = replace_for_await(TokenStream2::from(TokenStream::from(inner)));
48+
let syn::Block { stmts, .. } = syn::parse2(inner.clone())?;
4849

4950
let macro_ident = syn::Ident::new(
5051
&format!("stream_{}", count_bangs(inner.into())),
5152
Span::call_site(),
5253
);
5354

54-
AsyncStreamEnumHack { stmts, macro_ident }
55+
Ok(AsyncStreamEnumHack { stmts, macro_ident })
5556
}
5657
}
5758

@@ -153,7 +154,10 @@ pub fn async_stream_impl(input: TokenStream) -> TokenStream {
153154
let AsyncStreamEnumHack {
154155
macro_ident,
155156
mut stmts,
156-
} = AsyncStreamEnumHack::parse(input);
157+
} = match AsyncStreamEnumHack::parse(input) {
158+
Ok(x) => x,
159+
Err(e) => return e.to_compile_error().into(),
160+
};
157161

158162
let mut scrub = Scrub {
159163
is_xforming: true,
@@ -192,7 +196,10 @@ pub fn async_try_stream_impl(input: TokenStream) -> TokenStream {
192196
let AsyncStreamEnumHack {
193197
macro_ident,
194198
mut stmts,
195-
} = AsyncStreamEnumHack::parse(input);
199+
} = match AsyncStreamEnumHack::parse(input) {
200+
Ok(x) => x,
201+
Err(e) => return e.to_compile_error().into(),
202+
};
196203

197204
let mut scrub = Scrub {
198205
is_xforming: true,
@@ -245,3 +252,32 @@ fn count_bangs(input: TokenStream) -> usize {
245252

246253
count
247254
}
255+
256+
fn replace_for_await(input: TokenStream2) -> TokenStream2 {
257+
let mut input = input.into_iter().peekable();
258+
let mut tokens = Vec::new();
259+
260+
while let Some(token) = input.next() {
261+
match token {
262+
TokenTree2::Ident(ident) => {
263+
match input.peek() {
264+
Some(TokenTree2::Ident(next))
265+
if ident.to_string().as_str() == "for" && next == "await" =>
266+
{
267+
tokens.extend(quote!(#[#next]));
268+
let _ = input.next();
269+
}
270+
_ => {}
271+
}
272+
tokens.push(ident.into());
273+
}
274+
TokenTree2::Group(group) => {
275+
let stream = replace_for_await(group.stream());
276+
tokens.push(Group::new(group.delimiter(), stream).into());
277+
}
278+
_ => tokens.push(token),
279+
}
280+
}
281+
282+
tokens.into_iter().collect()
283+
}

async-stream/tests/for_await.rs

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,7 @@ async fn test() {
1212
};
1313

1414
let s = stream! {
15-
#[await]
16-
for x in s {
15+
for await x in s {
1716
yield x.to_owned() + "!";
1817
}
1918
};

0 commit comments

Comments
 (0)