pyo3_stub_gen_derive/gen_stub/
method.rs

1use super::{
2    arg::parse_args, escape_return_type, extract_documents, parse_pyo3_attrs, ArgInfo,
3    ArgsWithSignature, Attr, Signature,
4};
5
6use proc_macro2::TokenStream as TokenStream2;
7use quote::{quote, ToTokens, TokenStreamExt};
8use syn::{
9    Error, GenericArgument, ImplItemFn, PathArguments, Result, Type, TypePath, TypeReference,
10};
11
12#[derive(Debug, Clone, Copy, PartialEq)]
13pub enum MethodType {
14    Instance,
15    Static,
16    Class,
17    New,
18}
19
20#[derive(Debug)]
21pub struct MethodInfo {
22    name: String,
23    args: Vec<ArgInfo>,
24    sig: Option<Signature>,
25    r#return: Option<Type>,
26    doc: String,
27    r#type: MethodType,
28}
29
30fn replace_inner(ty: &mut Type, self_: &Type) {
31    match ty {
32        Type::Path(TypePath { path, .. }) => {
33            if let Some(last) = path.segments.last_mut() {
34                if let PathArguments::AngleBracketed(arg) = &mut last.arguments {
35                    for arg in &mut arg.args {
36                        if let GenericArgument::Type(ty) = arg {
37                            replace_inner(ty, self_);
38                        }
39                    }
40                }
41                if last.ident == "Self" {
42                    *ty = self_.clone();
43                }
44            }
45        }
46        Type::Reference(TypeReference { elem, .. }) => {
47            replace_inner(elem, self_);
48        }
49        _ => {}
50    }
51}
52
53impl MethodInfo {
54    pub fn replace_self(&mut self, self_: &Type) {
55        for arg in &mut self.args {
56            replace_inner(&mut arg.r#type, self_);
57        }
58        if let Some(ret) = self.r#return.as_mut() {
59            replace_inner(ret, self_);
60        }
61    }
62}
63
64impl TryFrom<ImplItemFn> for MethodInfo {
65    type Error = Error;
66    fn try_from(item: ImplItemFn) -> Result<Self> {
67        let ImplItemFn { attrs, sig, .. } = item;
68        let doc = extract_documents(&attrs).join("\n");
69        let attrs = parse_pyo3_attrs(&attrs)?;
70        let mut method_name = None;
71        let mut text_sig = Signature::overriding_operator(&sig);
72        let mut method_type = MethodType::Instance;
73        for attr in attrs {
74            match attr {
75                Attr::Name(name) => method_name = Some(name),
76                Attr::Signature(text_sig_) => text_sig = Some(text_sig_),
77                Attr::StaticMethod => method_type = MethodType::Static,
78                Attr::ClassMethod => method_type = MethodType::Class,
79                Attr::New => method_type = MethodType::New,
80                _ => {}
81            }
82        }
83        let name = if method_type == MethodType::New {
84            "__new__".to_string()
85        } else {
86            method_name.unwrap_or(sig.ident.to_string())
87        };
88        let r#return = escape_return_type(&sig.output);
89        Ok(MethodInfo {
90            name,
91            sig: text_sig,
92            args: parse_args(sig.inputs)?,
93            r#return,
94            doc,
95            r#type: method_type,
96        })
97    }
98}
99
100impl ToTokens for MethodInfo {
101    fn to_tokens(&self, tokens: &mut TokenStream2) {
102        let Self {
103            name,
104            r#return: ret,
105            args,
106            sig,
107            doc,
108            r#type,
109        } = self;
110        let args_with_sig = ArgsWithSignature { args, sig };
111        let ret_tt = if let Some(ret) = ret {
112            quote! { <#ret as pyo3_stub_gen::PyStubType>::type_output }
113        } else {
114            quote! { ::pyo3_stub_gen::type_info::no_return_type_output }
115        };
116        let type_tt = match r#type {
117            MethodType::Instance => quote! { ::pyo3_stub_gen::type_info::MethodType::Instance },
118            MethodType::Static => quote! { ::pyo3_stub_gen::type_info::MethodType::Static },
119            MethodType::Class => quote! { ::pyo3_stub_gen::type_info::MethodType::Class },
120            MethodType::New => quote! { ::pyo3_stub_gen::type_info::MethodType::New },
121        };
122        tokens.append_all(quote! {
123            ::pyo3_stub_gen::type_info::MethodInfo {
124                name: #name,
125                args: #args_with_sig,
126                r#return: #ret_tt,
127                doc: #doc,
128                r#type: #type_tt
129            }
130        })
131    }
132}
pyo3_stub_gen_derive/gen_stub/method.rs

pyo3_stub_gen_derive/gen_stub/
method.rs
