pyo3_stub_gen_derive/gen_stub/
pyclass_enum.rs

1use proc_macro2::TokenStream as TokenStream2;
2use quote::{quote, ToTokens, TokenStreamExt};
3use syn::{parse_quote, Error, ItemEnum, Result, Type};
4
5use super::{extract_documents, parse_pyo3_attrs, util::quote_option, Attr, StubType};
6
7pub struct PyEnumInfo {
8    pyclass_name: String,
9    enum_type: Type,
10    module: Option<String>,
11    variants: Vec<(String, String)>,
12    doc: String,
13}
14
15impl From<&PyEnumInfo> for StubType {
16    fn from(info: &PyEnumInfo) -> Self {
17        let PyEnumInfo {
18            pyclass_name,
19            module,
20            enum_type,
21            ..
22        } = info;
23        Self {
24            ty: enum_type.clone(),
25            name: pyclass_name.clone(),
26            module: module.clone(),
27        }
28    }
29}
30
31impl TryFrom<ItemEnum> for PyEnumInfo {
32    type Error = Error;
33    fn try_from(
34        ItemEnum {
35            variants,
36            attrs,
37            ident,
38            ..
39        }: ItemEnum,
40    ) -> Result<Self> {
41        let doc = extract_documents(&attrs).join("\n");
42        let mut pyclass_name = None;
43        let mut module = None;
44        let mut renaming_rule = None;
45        for attr in parse_pyo3_attrs(&attrs)? {
46            match attr {
47                Attr::Name(name) => pyclass_name = Some(name),
48                Attr::Module(name) => module = Some(name),
49                Attr::RenameAll(name) => renaming_rule = Some(name),
50                _ => {}
51            }
52        }
53        let struct_type = parse_quote!(#ident);
54        let pyclass_name = pyclass_name.unwrap_or_else(|| ident.to_string());
55        let variants = variants
56            .into_iter()
57            .map(|var| -> Result<(String, String)> {
58                let mut var_name = None;
59                for attr in parse_pyo3_attrs(&var.attrs)? {
60                    if let Attr::Name(name) = attr {
61                        var_name = Some(name);
62                    }
63                }
64                let mut var_name = var_name.unwrap_or_else(|| var.ident.to_string());
65                if let Some(renaming_rule) = renaming_rule {
66                    var_name = renaming_rule.apply(&var_name);
67                }
68                let var_doc = extract_documents(&var.attrs).join("\n");
69                Ok((var_name, var_doc))
70            })
71            .collect::<Result<Vec<(String, String)>>>()?;
72        Ok(Self {
73            doc,
74            enum_type: struct_type,
75            pyclass_name,
76            module,
77            variants,
78        })
79    }
80}
81
82impl ToTokens for PyEnumInfo {
83    fn to_tokens(&self, tokens: &mut TokenStream2) {
84        let Self {
85            pyclass_name,
86            enum_type,
87            variants,
88            doc,
89            module,
90        } = self;
91        let module = quote_option(module);
92        let variants: Vec<_> = variants
93            .iter()
94            .map(|(name, doc)| quote! {(#name,#doc)})
95            .collect();
96        tokens.append_all(quote! {
97            ::pyo3_stub_gen::type_info::PyEnumInfo {
98                pyclass_name: #pyclass_name,
99                enum_id: std::any::TypeId::of::<#enum_type>,
100                variants: &[ #(#variants),* ],
101                module: #module,
102                doc: #doc,
103            }
104        })
105    }
106}
pyo3_stub_gen_derive/gen_stub/pyclass_enum.rs

pyo3_stub_gen_derive/gen_stub/
pyclass_enum.rs
