pyo3_stub_gen/generate/
enum_.rs

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

pyo3_stub_gen/generate/
enum_.rs
