pyo3_stub_gen_derive/gen_stub/
pyfunction.rs

1use proc_macro2::TokenStream as TokenStream2;
2use quote::{quote, ToTokens, TokenStreamExt};
3use syn::{
4    parse::{Parse, ParseStream},
5    Error, ItemFn, Result, Type,
6};
7
8use super::{
9    escape_return_type, extract_documents, parse_args, parse_pyo3_attrs, quote_option, ArgInfo,
10    ArgsWithSignature, Attr, Signature,
11};
12
13pub struct PyFunctionInfo {
14    name: String,
15    args: Vec<ArgInfo>,
16    r#return: Option<Type>,
17    sig: Option<Signature>,
18    doc: String,
19    module: Option<String>,
20}
21
22struct ModuleAttr {
23    _module: syn::Ident,
24    _eq_token: syn::token::Eq,
25    name: syn::LitStr,
26}
27
28impl Parse for ModuleAttr {
29    fn parse(input: ParseStream) -> Result<Self> {
30        Ok(Self {
31            _module: input.parse()?,
32            _eq_token: input.parse()?,
33            name: input.parse()?,
34        })
35    }
36}
37
38impl PyFunctionInfo {
39    pub fn parse_attr(&mut self, attr: TokenStream2) -> Result<()> {
40        if attr.is_empty() {
41            return Ok(());
42        }
43        let attr: ModuleAttr = syn::parse2(attr)?;
44        self.module = Some(attr.name.value());
45        Ok(())
46    }
47}
48
49impl TryFrom<ItemFn> for PyFunctionInfo {
50    type Error = Error;
51    fn try_from(item: ItemFn) -> Result<Self> {
52        let doc = extract_documents(&item.attrs).join("\n");
53        let args = parse_args(item.sig.inputs)?;
54        let r#return = escape_return_type(&item.sig.output);
55        let mut name = None;
56        let mut sig = None;
57        for attr in parse_pyo3_attrs(&item.attrs)? {
58            match attr {
59                Attr::Name(function_name) => name = Some(function_name),
60                Attr::Signature(signature) => sig = Some(signature),
61                _ => {}
62            }
63        }
64        let name = name.unwrap_or_else(|| item.sig.ident.to_string());
65        Ok(Self {
66            args,
67            sig,
68            r#return,
69            name,
70            doc,
71            module: None,
72        })
73    }
74}
75
76impl ToTokens for PyFunctionInfo {
77    fn to_tokens(&self, tokens: &mut TokenStream2) {
78        let Self {
79            args,
80            r#return: ret,
81            name,
82            doc,
83            sig,
84            module,
85        } = self;
86        let ret_tt = if let Some(ret) = ret {
87            quote! { <#ret as pyo3_stub_gen::PyStubType>::type_output }
88        } else {
89            quote! { ::pyo3_stub_gen::type_info::no_return_type_output }
90        };
91        // let sig_tt = quote_option(sig);
92        let module_tt = quote_option(module);
93        let args_with_sig = ArgsWithSignature { args, sig };
94        tokens.append_all(quote! {
95            ::pyo3_stub_gen::type_info::PyFunctionInfo {
96                name: #name,
97                args: #args_with_sig,
98                r#return: #ret_tt,
99                doc: #doc,
100                module: #module_tt,
101            }
102        })
103    }
104}
pyo3_stub_gen_derive/gen_stub/pyfunction.rs

pyo3_stub_gen_derive/gen_stub/
pyfunction.rs
