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(tokens: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}
105