timeout_macro_parse/
lib.rs

1extern crate proc_macro;
2
3use crate::compile_error::to_compile_error;
4use crate::inject::{try_inject, Injector};
5use crate::parse_attr::{parse_attr, ValidOpts};
6#[cfg(not(feature = "test"))]
7use proc_macro::{Delimiter, Group, Ident, Punct, Spacing, Span, TokenStream, TokenTree};
8#[cfg(feature = "test")]
9use proc_macro2::{Delimiter, Group, Ident, Punct, Spacing, Span, TokenStream, TokenTree};
10use std::fmt::Display;
11
12mod compile_error;
13pub mod inject;
14pub mod parse_attr;
15pub mod parse_duration;
16
17struct TokioTimeoutInjector(ValidOpts);
18
19pub(crate) enum Error {
20    Parse(Span, String),
21    ParseSpanMissing(String),
22}
23
24impl Error {
25    pub fn missing_span(msg: String) -> Self {
26        Self::ParseSpanMissing(msg)
27    }
28
29    pub fn with_span<T: Display>(span: Span, msg: T) -> Self {
30        Self::Parse(span, msg.to_string())
31    }
32
33    pub fn with_span_if_missing(self, span: Span) -> Self {
34        match self {
35            Self::Parse(_, _) => self,
36            Self::ParseSpanMissing(e) => Self::Parse(span, e),
37        }
38    }
39
40    pub fn into_token_stream_with_fallback_span(self, span: Span) -> TokenStream {
41        match self {
42            Self::Parse(p, msg) => to_compile_error(&msg, p),
43            Self::ParseSpanMissing(e) => to_compile_error(&e, span),
44        }
45    }
46}
47
48pub(crate) type Result<T> = core::result::Result<T, Error>;
49
50impl Injector for TokioTimeoutInjector {
51    fn inject(self, fn_name: &str, inner_code: TokenStream) -> TokenStream {
52        let err_disp = self.0.duration.to_error_display(fn_name);
53        let on_timeout = self.0.on_error.into_token_stream(&err_disp);
54        let dur = self.0.duration.into_token_stream();
55        let mut inner = TokenStream::new();
56        let span = Span::call_site();
57        let mut timeout_args = TokenStream::new();
58        timeout_args.extend(dur);
59        timeout_args.extend([
60            TokenTree::Punct(Punct::new(',', Spacing::Alone)),
61            TokenTree::Ident(Ident::new("async", span)),
62            TokenTree::Group(Group::new(Delimiter::Brace, inner_code)),
63        ]);
64        let mut match_body = TokenStream::new();
65        let mut value_group = TokenStream::new();
66        value_group.extend([TokenTree::Ident(Ident::new("v", span))]);
67        let mut err_group = TokenStream::new();
68        err_group.extend([TokenTree::Ident(Ident::new("e", span))]);
69        match_body.extend([
70            TokenTree::Ident(Ident::new("Ok", span)),
71            TokenTree::Group(Group::new(Delimiter::Parenthesis, value_group)),
72            TokenTree::Punct(Punct::new('=', Spacing::Joint)),
73            TokenTree::Punct(Punct::new('>', Spacing::Alone)),
74            TokenTree::Ident(Ident::new("v", span)),
75            TokenTree::Punct(Punct::new(',', Spacing::Alone)),
76            TokenTree::Ident(Ident::new("Err", span)),
77            TokenTree::Group(Group::new(Delimiter::Parenthesis, err_group)),
78            TokenTree::Punct(Punct::new('=', Spacing::Joint)),
79            TokenTree::Punct(Punct::new('>', Spacing::Alone)),
80        ]);
81        match_body.extend(on_timeout);
82        inner.extend([
83            TokenTree::Ident(Ident::new("match", span)),
84            TokenTree::Ident(Ident::new("tokio", span)),
85            TokenTree::Punct(Punct::new(':', Spacing::Joint)),
86            TokenTree::Punct(Punct::new(':', Spacing::Alone)),
87            TokenTree::Ident(Ident::new("time", span)),
88            TokenTree::Punct(Punct::new(':', Spacing::Joint)),
89            TokenTree::Punct(Punct::new(':', Spacing::Alone)),
90            TokenTree::Ident(Ident::new("timeout", span)),
91            TokenTree::Group(Group::new(Delimiter::Parenthesis, timeout_args)),
92            TokenTree::Punct(Punct::new('.', Spacing::Alone)),
93            TokenTree::Ident(Ident::new("await", span)),
94            TokenTree::Group(Group::new(Delimiter::Brace, match_body)),
95        ]);
96        TokenStream::from(TokenTree::Group(Group::new(Delimiter::Brace, inner)))
97    }
98}
99
100pub fn tokio_timeout(attr: TokenStream, item: TokenStream) -> TokenStream {
101    let validated = match parse_attr(attr) {
102        Ok(o) => o,
103        Err(e) => {
104            return e.into_token_stream_with_fallback_span(Span::call_site());
105        }
106    };
107    try_inject(TokioTimeoutInjector(validated), item)
108        .unwrap_or_else(|e| e.into_token_stream_with_fallback_span(Span::call_site()))
109}
timeout_macro_parse/lib.rs

timeout_macro_parse/
lib.rs
