pyo3_stub_gen/generate/
class.rs

1use crate::{generate::*, type_info::*, TypeInfo};
2use std::fmt;
3
4/// Definition of a Python class.
5#[derive(Debug, Clone, PartialEq)]
6pub struct ClassDef {
7    pub name: &'static str,
8    pub doc: &'static str,
9    pub members: Vec<MemberDef>,
10    pub methods: Vec<MethodDef>,
11    pub bases: Vec<TypeInfo>,
12}
13
14impl Import for ClassDef {
15    fn import(&self) -> HashSet<ModuleRef> {
16        let mut import = HashSet::new();
17        for base in &self.bases {
18            import.extend(base.import.clone());
19        }
20        for member in &self.members {
21            import.extend(member.import());
22        }
23        for method in &self.methods {
24            import.extend(method.import());
25        }
26        import
27    }
28}
29
30impl From<&PyClassInfo> for ClassDef {
31    fn from(info: &PyClassInfo) -> Self {
32        // Since there are multiple `#[pymethods]` for a single class, we need to merge them.
33        // This is only an initializer. See `StubInfo::gather` for the actual merging.
34        Self {
35            name: info.pyclass_name,
36            doc: info.doc,
37            members: info.members.iter().map(MemberDef::from).collect(),
38            methods: Vec::new(),
39            bases: info.bases.iter().map(|f| f()).collect(),
40        }
41    }
42}
43
44impl fmt::Display for ClassDef {
45    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
46        let bases = self
47            .bases
48            .iter()
49            .map(|i| i.name.clone())
50            .reduce(|acc, path| format!("{acc}, {path}"))
51            .map(|bases| format!("({bases})"))
52            .unwrap_or_default();
53        writeln!(f, "class {}{}:", self.name, bases)?;
54        let indent = indent();
55        let doc = self.doc.trim();
56        docstring::write_docstring(f, doc, indent)?;
57        for member in &self.members {
58            member.fmt(f)?;
59        }
60        for method in &self.methods {
61            method.fmt(f)?;
62        }
63        if self.members.is_empty() && self.methods.is_empty() {
64            writeln!(f, "{indent}...")?;
65        }
66        writeln!(f)?;
67        Ok(())
68    }
69}