oscript_derive/
lib.rs

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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
#![feature(proc_macro_diagnostic)]
extern crate proc_macro;
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, FnArg, ItemFn, Pat, PatType, Type};


#[proc_macro_attribute]
pub fn oscript_async_main(attr: TokenStream, item: TokenStream) -> TokenStream {
    oscript_main_internal(attr, item, true)
}


#[proc_macro_attribute]
pub fn oscript_main(attr: TokenStream, item: TokenStream) -> TokenStream {
    oscript_main_internal(attr, item, false)
}

fn oscript_main_internal(attr: TokenStream, item: TokenStream,uses_async_macro:bool) -> TokenStream {
    
    let mut use_tokio = uses_async_macro;

    let parser = syn::meta::parser(|meta| {
        if meta.path.is_ident("use_tokio") {
            use_tokio = true;
            Ok(())
        } else {
            Err(meta.error("unsupported attribute for [oscript..] on your main method.."))
        }
    });
    
    parse_macro_input!(attr with parser);

    let input = parse_macro_input!(item as ItemFn);

    let is_marked_with_async = input.sig.asyncness.is_some();

    if is_marked_with_async && use_tokio == false {
        panic!("Your main method is marked as async but you have not enabled tokio... Try using: #[oscript_async_main]")
    }
    else if !is_marked_with_async && use_tokio  {
        proc_macro::Span::call_site().warning("#[oscript_async_main] used on non-async main will implicitly make it async! Consider marking it for clarity.").emit();
    }

    // Extract function name and arguments
    let fn_name = &input.sig.ident;
    if fn_name != "main" {
        panic!("Only the `main` function can be annotated with #[oscript_main]");
    }
    let fn_args = &input.sig.inputs;

    let struct_name = format_ident!("O{}Args", fn_name);

    let mut struct_fields = Vec::new();
    let mut fn_arg_names = Vec::new();
    let mut fn_arg_conversions = Vec::new();

    for arg in fn_args.iter() {
        if let FnArg::Typed(PatType { pat, ty, .. }) = arg {
            if let Pat::Ident(pat_ident) = &**pat {
                let arg_name = &pat_ident.ident;
                match &**ty {
                    Type::Reference(_) => {
                        // If the argument is a reference we will just use string.. not sure what else to do atm
                        struct_fields.push(quote! {
                            #[clap(long)]
                            pub #arg_name: String
                        });
                        fn_arg_names.push(arg_name.clone());
                        fn_arg_conversions.push(quote! {
                            let #arg_name = args.#arg_name.as_str();
                        });
                    }
                    _ => {
                        struct_fields.push(quote! {
                            #[clap(long)]
                            pub #arg_name: #ty
                        });
                        fn_arg_names.push(arg_name.clone());
                        fn_arg_conversions.push(quote! {
                            let #arg_name = args.#arg_name.clone();
                        });
                    }
                }
            }
        }
    }

    // Add a custom flag for generating shell script IntelliSense
    struct_fields.push(quote! {
        #[clap(long, hide = false)]
        pub generate_completion: Option<String>
    });

    // Handle IntelliSense flag logic so we can help people generate script completion :D
    let completion_logic = quote! {
        use clap_complete::Shell;
        if let Some(shell) = &args.generate_completion {
            let mut app = <#struct_name as clap::CommandFactory>::command();
            let shell = shell.parse::<Shell>().expect("Invalid shell type");
            clap_complete::generate(shell, &mut app, env!("CARGO_PKG_NAME"), &mut std::io::stdout());
            return;
        }
    };

    // Remove arguments from the function since main cant have them
    let fn_block = &input.block;
    let visibility = &input.vis;

    let clap_quote = quote! {
        #[derive(clap::Parser, Debug)]
        pub struct #struct_name {
            #(#struct_fields),*
        }
    };

    // probably just join these if there is no more diff than this..
    let output = if use_tokio {
        quote! {
            use clap_complete::Shell;
            #clap_quote
            #[tokio::main]
            #visibility async fn #fn_name() {
                let args = <#struct_name as clap::Parser>::parse();
                #completion_logic
                #(#fn_arg_conversions)*
                #fn_block
            }
        }
    } else {
        quote! {
            #clap_quote
            #visibility fn #fn_name() {
                let args = <#struct_name as clap::Parser>::parse();
                #completion_logic
                #(#fn_arg_conversions)*
                #fn_block
            }
        }
    };

    TokenStream::from(output)
}