pyo3_stub_gen/generate/
enum_.rs

1use crate::{generate::*, type_info::*};
2use std::fmt;
3
4/// Definition of a Python enum.
5#[derive(Debug, Clone, PartialEq)]
6pub struct EnumDef {
7    pub name: &'static str,
8    pub doc: &'static str,
9    pub variants: &'static [(&'static str, &'static str)],
10    pub methods: Vec<MethodDef>,
11    pub attrs: Vec<MemberDef>,
12    pub getters: Vec<MemberDef>,
13    pub setters: Vec<MemberDef>,
14}
15
16impl From<&PyEnumInfo> for EnumDef {
17    fn from(info: &PyEnumInfo) -> Self {
18        Self {
19            name: info.pyclass_name,
20            doc: info.doc,
21            variants: info.variants,
22            methods: Vec::new(),
23            attrs: Vec::new(),
24            getters: Vec::new(),
25            setters: Vec::new(),
26        }
27    }
28}
29
30impl Import for EnumDef {
31    fn import(&self) -> HashSet<ImportRef> {
32        let mut import = HashSet::new();
33        // for @typing.final
34        import.insert("typing".into());
35        // for Enum base class
36        import.insert("enum".into());
37        for method in &self.methods {
38            import.extend(method.import());
39        }
40        for attr in &self.attrs {
41            import.extend(attr.import());
42        }
43        for getter in &self.getters {
44            import.extend(getter.import());
45        }
46        for setter in &self.setters {
47            import.extend(setter.import());
48        }
49        import
50    }
51}
52
53impl fmt::Display for EnumDef {
54    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
55        writeln!(f, "@typing.final")?;
56        writeln!(f, "class {}(enum.Enum):", self.name)?;
57        let indent = indent();
58        docstring::write_docstring(f, self.doc, indent)?;
59        for (variant, variant_doc) in self.variants {
60            writeln!(f, "{indent}{variant} = ...")?;
61            docstring::write_docstring(f, variant_doc, indent)?;
62        }
63        if !(self.attrs.is_empty()
64            && self.getters.is_empty()
65            && self.setters.is_empty()
66            && self.methods.is_empty())
67        {
68            writeln!(f)?;
69            for attr in &self.attrs {
70                attr.fmt(f)?;
71            }
72            for getter in &self.getters {
73                GetterDisplay(getter).fmt(f)?;
74            }
75            for setter in &self.setters {
76                SetterDisplay(setter).fmt(f)?;
77            }
78            for methods in &self.methods {
79                methods.fmt(f)?;
80            }
81        }
82        writeln!(f)?;
83        Ok(())
84    }
85}
pyo3_stub_gen/generate/enum_.rs

pyo3_stub_gen/generate/
enum_.rs
