easy_macros_add_code/
lib.rs

1use anyhow::bail;
2use proc_macro::TokenStream;
3use quote::ToTokens;
4use syn::{
5    Block, Ident, Token,
6    parse::{Parse, ParseStream},
7};
8
9#[derive(Default)]
10struct AddCodeArgs {
11    before: Option<Vec<syn::Stmt>>,
12    after: Option<Vec<syn::Stmt>>,
13}
14
15impl Parse for AddCodeArgs {
16    fn parse(input: ParseStream) -> syn::Result<Self> {
17        let mut args = AddCodeArgs::default();
18
19        while !input.is_empty() {
20            let name: Ident = input.parse()?;
21            input.parse::<Token![=]>()?;
22
23            let block: Block = input.parse()?;
24            let stmts = block.stmts;
25
26            match name.to_string().as_str() {
27                "before" => {
28                    if args.before.is_some() {
29                        return Err(syn::Error::new(name.span(), "duplicate `before`"));
30                    }
31                    args.before = Some(stmts);
32                }
33                "after" => {
34                    if args.after.is_some() {
35                        return Err(syn::Error::new(name.span(), "duplicate `after`"));
36                    }
37                    args.after = Some(stmts);
38                }
39                _ => {
40                    return Err(syn::Error::new(name.span(), "expected `before` or `after`"));
41                }
42            }
43
44            if input.peek(Token![,]) {
45                input.parse::<Token![,]>()?;
46            }
47        }
48
49        if args.before.is_none() && args.after.is_none() {
50            return Err(syn::Error::new(input.span(), "missing `before` or `after`"));
51        }
52
53        Ok(args)
54    }
55}
56
57fn inject_into_block(block: &mut Block, args: AddCodeArgs) {
58    if let Some(before) = args.before {
59        block.stmts.splice(0..0, before);
60    }
61
62    if let Some(after) = args.after {
63        block.stmts.extend(after);
64    }
65}
66
67fn add_code_base(attr: TokenStream, item: TokenStream) -> anyhow::Result<TokenStream> {
68    let args = helpers::parse_macro_input!(attr as AddCodeArgs);
69
70    let item_ts: proc_macro2::TokenStream = item.clone().into();
71
72    if let Ok(mut item_fn) = syn::parse2::<syn::ImplItemFn>(item_ts) {
73        inject_into_block(&mut item_fn.block, args);
74        return Ok(item_fn.to_token_stream().into());
75    }
76
77    bail!("#[add_code] can only be used on functions or impl methods");
78}
79
80#[proc_macro_attribute]
81#[anyhow_result::anyhow_result]
82/// Injects code at the beginning and/or end of a function-like item.
83///
84/// This is handy for keeping [docify](https://crates.io/crates/docify)-generated examples clean while still
85/// inserting setup/teardown or assertions for tests.
86///
87/// # Syntax
88/// ```rust,ignore
89/// #[add_code(before = { /* statements */ })]
90/// #[add_code(after = { /* statements */ })]
91/// #[add_code(before = { ... }, after = { ... })]
92/// fn demo() { /* ... */ }
93/// ```
94///
95/// The statements inside the braces are inserted without the outer `{}`.
96pub fn add_code(attr: TokenStream, item: TokenStream) -> anyhow::Result<TokenStream> {
97    add_code_base(attr, item)
98}
99#[proc_macro_attribute]
100#[anyhow_result::anyhow_result]
101#[doc(hidden)]
102pub fn add_code_debug(attr: TokenStream, item: TokenStream) -> anyhow::Result<TokenStream> {
103    panic!("{}", add_code_base(attr, item)?)
104}
easy_macros_add_code/lib.rs

easy_macros_add_code/
lib.rs
