pollster-macro 0.4.0

Proc-macro crate for pollster
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

#![doc = include_str!("../README.md")]

use std::iter::FromIterator;
use std::str::FromStr;

use proc_macro2::TokenStream;
use quote::ToTokens;
use syn::spanned::Spanned;
use syn::{Error, Expr, ExprLit, ExprPath, ItemFn, Lit, MetaNameValue, Result};

/// Uses [`pollster::block_on`] to enable `async fn main() {}`.
///
/// # Example
///
/// ```
/// #[pollster::main]
/// async fn main() {
///     let my_fut = async {};
///
///     my_fut.await;
/// }
/// ```
///
/// [`pollster::block_on`]: https://docs.rs/pollster/0.3.0/pollster/fn.block_on.html
#[proc_macro_attribute]
pub fn main(
    attr: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let item = TokenStream::from(item);
    let backup = item.clone();

    match common(attr.into(), item) {
        Ok(output) => output.into_token_stream().into(),
        Err(error) => TokenStream::from_iter([error.into_compile_error(), backup]).into(),
    }
}

/// Uses [`pollster::block_on`] to enable `async` on test functions.
///
/// # Example
///
/// ```ignore
/// #[pollster::test]
/// async fn main() {
///     let my_fut = async {};
///
///     my_fut.await;
/// }
/// ```
///
/// [`pollster::block_on`]: https://docs.rs/pollster/0.3.0/pollster/fn.block_on.html
#[proc_macro_attribute]
pub fn test(
    attr: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let item = TokenStream::from(item);
    let backup = item.clone();

    match test_internal(attr.into(), item) {
        Ok(output) => output.into_token_stream().into(),
        Err(error) => TokenStream::from_iter([error.into_compile_error(), backup]).into(),
    }
}

fn test_internal(attr: TokenStream, item: TokenStream) -> Result<ItemFn> {
    let mut item = common(attr, item)?;
    item.attrs.push(syn::parse_quote! { #[test] });

    Ok(item)
}

fn common(attr: TokenStream, item: TokenStream) -> Result<ItemFn> {
    let mut item: ItemFn = syn::parse2(item)?;

    if item.sig.asyncness.is_some() {
        item.sig.asyncness = None;
    } else {
        return Err(Error::new_spanned(item, "expected function to be async"));
    }

    let path = if attr.is_empty() {
        quote::quote! { ::pollster }
    } else {
        let attr: MetaNameValue = syn::parse2(attr)?;

        if attr.path.is_ident("crate") {
            match attr.value {
                Expr::Lit(ExprLit {
                    attrs,
                    lit: Lit::Str(str),
                }) if attrs.is_empty() => TokenStream::from_str(&str.value())?,
                Expr::Path(ExprPath {
                    attrs,
                    qself: None,
                    path,
                }) if attrs.is_empty() => path.to_token_stream(),
                _ => {
                    return Err(Error::new_spanned(
                        attr.value,
                        "expected valid path, e.g. `::package_name`",
                    ))
                }
            }
        } else {
            return Err(Error::new_spanned(attr.path, "expected `crate`"));
        }
    };

    let span = item.span();
    let block = item.block;
    item.block = syn::parse_quote_spanned! {
        span =>
        {
            #path::block_on(async {
                #block
            })
        }
    };

    Ok(item)
}