pyo3_stub_gen_derive/gen_stub/
pymethods.rs

1use proc_macro2::TokenStream as TokenStream2;
2use quote::{quote, ToTokens, TokenStreamExt};
3use syn::{Error, FnArg, ImplItem, ItemImpl, Result, Type};
4
5use super::{attr::parse_gen_stub_skip, MemberInfo, MethodInfo};
6
7#[derive(Debug)]
8pub struct PyMethodsInfo {
9    struct_id: Type,
10    attrs: Vec<MemberInfo>,
11    getters: Vec<MemberInfo>,
12    setters: Vec<MemberInfo>,
13    methods: Vec<MethodInfo>,
14}
15
16impl TryFrom<ItemImpl> for PyMethodsInfo {
17    type Error = Error;
18    fn try_from(item: ItemImpl) -> Result<Self> {
19        let struct_id = *item.self_ty.clone();
20        let mut attrs = Vec::new();
21        let mut getters = Vec::new();
22        let mut setters = Vec::new();
23        let mut methods = Vec::new();
24        for inner in item.items.into_iter() {
25            match inner {
26                ImplItem::Const(item_const) => {
27                    if parse_gen_stub_skip(&item_const.attrs)? {
28                        continue;
29                    }
30                    if MemberInfo::is_classattr(&item_const.attrs)? {
31                        attrs.push(MemberInfo::new_classattr_const(item_const)?);
32                    }
33                }
34                ImplItem::Fn(item_fn) => {
35                    if parse_gen_stub_skip(&item_fn.attrs)? {
36                        continue;
37                    }
38                    if MemberInfo::is_getter(&item_fn.attrs)? {
39                        getters.push(MemberInfo::new_getter(item_fn)?);
40                        continue;
41                    }
42                    if MemberInfo::is_setter(&item_fn.attrs)? {
43                        setters.push(MemberInfo::new_setter(item_fn)?);
44                        continue;
45                    }
46                    if MemberInfo::is_classattr(&item_fn.attrs)? {
47                        attrs.push(MemberInfo::new_classattr_fn(item_fn)?);
48                        continue;
49                    }
50                    let mut method = MethodInfo::try_from(item_fn)?;
51                    method.replace_self(&item.self_ty);
52                    methods.push(method);
53                }
54                _ => continue,
55            }
56        }
57        Ok(Self {
58            struct_id,
59            attrs,
60            getters,
61            setters,
62            methods,
63        })
64    }
65}
66
67impl ToTokens for PyMethodsInfo {
68    fn to_tokens(&self, tokens: &mut TokenStream2) {
69        let Self {
70            struct_id,
71            attrs,
72            getters,
73            setters,
74            methods,
75        } = self;
76        tokens.append_all(quote! {
77            ::pyo3_stub_gen::type_info::PyMethodsInfo {
78                struct_id: std::any::TypeId::of::<#struct_id>,
79                attrs: &[ #(#attrs),* ],
80                getters: &[ #(#getters),* ],
81                setters: &[ #(#setters),* ],
82                methods: &[ #(#methods),* ],
83            }
84        })
85    }
86}
87
88// `#[gen_stub(xxx)]` is not a valid proc_macro_attribute
89// it's only designed to receive user's setting.
90// We need to remove all `#[gen_stub(xxx)]` before print the item_impl back
91pub fn prune_attrs(item_impl: &mut ItemImpl) {
92    super::attr::prune_attrs(&mut item_impl.attrs);
93    for inner in item_impl.items.iter_mut() {
94        if let ImplItem::Fn(item_fn) = inner {
95            super::attr::prune_attrs(&mut item_fn.attrs);
96            for arg in item_fn.sig.inputs.iter_mut() {
97                if let FnArg::Typed(ref mut pat_type) = arg {
98                    super::attr::prune_attrs(&mut pat_type.attrs);
99                }
100            }
101        }
102    }
103}