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 module: Option<&'static str>,
9    pub doc: &'static str,
10    pub variants: &'static [(&'static str, &'static str)],
11    pub methods: Vec<MethodDef>,
12    pub attrs: Vec<MemberDef>,
13    pub getters: Vec<MemberDef>,
14    pub setters: Vec<MemberDef>,
15}
16
17impl From<&PyEnumInfo> for EnumDef {
18    fn from(info: &PyEnumInfo) -> Self {
19        Self {
20            name: info.pyclass_name,
21            module: info.module,
22            doc: info.doc,
23            variants: info.variants,
24            methods: Vec::new(),
25            attrs: Vec::new(),
26            getters: Vec::new(),
27            setters: Vec::new(),
28        }
29    }
30}
31
32impl Import for EnumDef {
33    fn import(&self) -> HashSet<ImportRef> {
34        let mut import = HashSet::new();
35        // for @typing.final
36        import.insert("typing".into());
37        // for Enum base class
38        import.insert("enum".into());
39        for method in &self.methods {
40            import.extend(method.import());
41        }
42        for attr in &self.attrs {
43            import.extend(attr.import());
44        }
45        for getter in &self.getters {
46            import.extend(getter.import());
47        }
48        for setter in &self.setters {
49            import.extend(setter.import());
50        }
51        import
52    }
53}
54
55impl fmt::Display for EnumDef {
56    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
57        writeln!(f, "@typing.final")?;
58        writeln!(f, "class {}(enum.Enum):", self.name)?;
59        let indent = indent();
60        docstring::write_docstring(f, self.doc, indent)?;
61        for (variant, variant_doc) in self.variants {
62            writeln!(f, "{indent}{variant} = ...")?;
63            docstring::write_docstring(f, variant_doc, indent)?;
64        }
65        if !(self.attrs.is_empty()
66            && self.getters.is_empty()
67            && self.setters.is_empty()
68            && self.methods.is_empty())
69        {
70            writeln!(f)?;
71            for attr in &self.attrs {
72                attr.fmt(f)?;
73            }
74            for getter in &self.getters {
75                write!(
76                    f,
77                    "{}",
78                    GetterDisplay {
79                        member: getter,
80                        target_module: self.module.unwrap_or(self.name)
81                    }
82                )?;
83            }
84            for setter in &self.setters {
85                write!(
86                    f,
87                    "{}",
88                    SetterDisplay {
89                        member: setter,
90                        target_module: self.module.unwrap_or(self.name)
91                    }
92                )?;
93            }
94            for methods in &self.methods {
95                methods.fmt(f)?;
96            }
97        }
98        writeln!(f)?;
99        Ok(())
100    }
101}
102
103impl EnumDef {
104    /// Format enum with module-qualified type names
105    ///
106    /// This method uses the target module context to qualify type identifiers
107    /// within compound type expressions based on their source modules.
108    /// Note: Enums currently don't have TypeInfo in their base classes, so this
109    /// mostly delegates to Display, but is provided for API consistency.
110    pub fn fmt_for_module(&self, target_module: &str, f: &mut fmt::Formatter) -> fmt::Result {
111        writeln!(f, "@typing.final")?;
112        writeln!(f, "class {}(enum.Enum):", self.name)?;
113        let indent = indent();
114        docstring::write_docstring(f, self.doc, indent)?;
115        for (variant, variant_doc) in self.variants {
116            writeln!(f, "{indent}{variant} = ...")?;
117            docstring::write_docstring(f, variant_doc, indent)?;
118        }
119        if !(self.attrs.is_empty()
120            && self.getters.is_empty()
121            && self.setters.is_empty()
122            && self.methods.is_empty())
123        {
124            writeln!(f)?;
125            for attr in &self.attrs {
126                attr.fmt_for_module(target_module, f, indent)?;
127            }
128            for getter in &self.getters {
129                write!(
130                    f,
131                    "{}",
132                    GetterDisplay {
133                        member: getter,
134                        target_module
135                    }
136                )?;
137            }
138            for setter in &self.setters {
139                write!(
140                    f,
141                    "{}",
142                    SetterDisplay {
143                        member: setter,
144                        target_module
145                    }
146                )?;
147            }
148            for methods in &self.methods {
149                methods.fmt_for_module(target_module, f, indent)?;
150            }
151        }
152        writeln!(f)?;
153        Ok(())
154    }
155}
pyo3_stub_gen/generate/enum_.rs

pyo3_stub_gen/generate/
enum_.rs
