timeout_macro_parse/
lib.rs

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

timeout_macro_parse/
lib.rs
