timeout-macro-parse 0.2.0

A companion-crate for the tokio-timeout proc-macro-lib
Documentation
#![warn(clippy::pedantic)]
extern crate proc_macro;

use crate::compile_error::to_compile_error;
use crate::inject::{try_inject, Injector};
use crate::parse_attr::{parse_attr, ValidOpts};
#[cfg(not(feature = "test"))]
use proc_macro::{Delimiter, Group, Ident, Punct, Spacing, Span, TokenStream, TokenTree};
#[cfg(feature = "test")]
use proc_macro2::{Delimiter, Group, Ident, Punct, Spacing, Span, TokenStream, TokenTree};
use std::fmt::Display;

mod compile_error;
mod inject;
mod parse_attr;
mod parse_duration;

struct TokioTimeoutInjector(ValidOpts);

pub(crate) enum Error {
    Parse(Span, String),
    ParseSpanMissing(String),
}

impl Error {
    pub(crate) fn missing_span(msg: String) -> Self {
        Self::ParseSpanMissing(msg)
    }

    pub(crate) fn with_span<T: Display>(span: Span, msg: T) -> Self {
        Self::Parse(span, msg.to_string())
    }

    pub(crate) fn with_span_if_missing(self, span: Span) -> Self {
        match self {
            Self::Parse(_, _) => self,
            Self::ParseSpanMissing(e) => Self::Parse(span, e),
        }
    }

    pub(crate) fn into_token_stream_with_fallback_span(self, span: Span) -> TokenStream {
        match self {
            Self::Parse(p, msg) => to_compile_error(&msg, p),
            Self::ParseSpanMissing(e) => to_compile_error(&e, span),
        }
    }
}

pub(crate) type Result<T> = core::result::Result<T, Error>;

impl Injector for TokioTimeoutInjector {
    fn inject(self, fn_name: &str, inner_code: TokenStream) -> TokenStream {
        let err_disp = self.0.duration.to_error_display(fn_name);
        let on_timeout = self.0.on_error.into_token_stream(&err_disp);
        let dur = self.0.duration.into_token_stream();
        let mut inner = TokenStream::new();
        let span = Span::call_site();
        let mut timeout_args = TokenStream::new();
        timeout_args.extend(dur);
        timeout_args.extend([
            TokenTree::Punct(Punct::new(',', Spacing::Alone)),
            TokenTree::Ident(Ident::new("async", span)),
            TokenTree::Group(Group::new(Delimiter::Brace, inner_code)),
        ]);
        let mut match_body = TokenStream::new();
        let mut value_group = TokenStream::new();
        value_group.extend([TokenTree::Ident(Ident::new("v", span))]);
        let mut err_group = TokenStream::new();
        err_group.extend([TokenTree::Ident(Ident::new("e", span))]);
        match_body.extend([
            TokenTree::Ident(Ident::new("Ok", span)),
            TokenTree::Group(Group::new(Delimiter::Parenthesis, value_group)),
            TokenTree::Punct(Punct::new('=', Spacing::Joint)),
            TokenTree::Punct(Punct::new('>', Spacing::Alone)),
            TokenTree::Ident(Ident::new("v", span)),
            TokenTree::Punct(Punct::new(',', Spacing::Alone)),
            TokenTree::Ident(Ident::new("Err", span)),
            TokenTree::Group(Group::new(Delimiter::Parenthesis, err_group)),
            TokenTree::Punct(Punct::new('=', Spacing::Joint)),
            TokenTree::Punct(Punct::new('>', Spacing::Alone)),
        ]);
        match_body.extend(on_timeout);
        inner.extend([
            TokenTree::Ident(Ident::new("match", span)),
            TokenTree::Ident(Ident::new("tokio", span)),
            TokenTree::Punct(Punct::new(':', Spacing::Joint)),
            TokenTree::Punct(Punct::new(':', Spacing::Alone)),
            TokenTree::Ident(Ident::new("time", span)),
            TokenTree::Punct(Punct::new(':', Spacing::Joint)),
            TokenTree::Punct(Punct::new(':', Spacing::Alone)),
            TokenTree::Ident(Ident::new("timeout", span)),
            TokenTree::Group(Group::new(Delimiter::Parenthesis, timeout_args)),
            TokenTree::Punct(Punct::new('.', Spacing::Alone)),
            TokenTree::Ident(Ident::new("await", span)),
            TokenTree::Group(Group::new(Delimiter::Brace, match_body)),
        ]);
        TokenStream::from(TokenTree::Group(Group::new(Delimiter::Brace, inner)))
    }
}

#[must_use]
pub fn tokio_timeout(attr: TokenStream, item: TokenStream) -> TokenStream {
    let validated = match parse_attr(attr) {
        Ok(o) => o,
        Err(e) => {
            return e.into_token_stream_with_fallback_span(Span::call_site());
        }
    };
    try_inject(TokioTimeoutInjector(validated), item)
        .unwrap_or_else(|e| e.into_token_stream_with_fallback_span(Span::call_site()))
}