pyo3_stub_gen/generate/
method.rs

1use crate::{generate::*, type_info::*, TypeInfo};
2use std::{collections::HashSet, fmt};
3
4pub use crate::type_info::MethodType;
5
6/// Definition of a class method.
7#[derive(Debug, Clone, PartialEq)]
8pub struct MethodDef {
9    pub name: &'static str,
10    pub args: Vec<Arg>,
11    pub r#return: TypeInfo,
12    pub doc: &'static str,
13    pub r#type: MethodType,
14}
15
16impl Import for MethodDef {
17    fn import(&self) -> HashSet<ModuleRef> {
18        let mut import = self.r#return.import.clone();
19        for arg in &self.args {
20            import.extend(arg.import().into_iter());
21        }
22        import
23    }
24}
25
26impl From<&MethodInfo> for MethodDef {
27    fn from(info: &MethodInfo) -> Self {
28        Self {
29            name: info.name,
30            args: info.args.iter().map(Arg::from).collect(),
31            r#return: (info.r#return)(),
32            doc: info.doc,
33            r#type: info.r#type,
34        }
35    }
36}
37
38impl fmt::Display for MethodDef {
39    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
40        let indent = indent();
41        let mut needs_comma = false;
42        match self.r#type {
43            MethodType::Static => {
44                writeln!(f, "{indent}@staticmethod")?;
45                write!(f, "{indent}def {}(", self.name)?;
46            }
47            MethodType::Class | MethodType::New => {
48                if self.r#type == MethodType::Class {
49                    // new is a classmethod without the decorator
50                    writeln!(f, "{indent}@classmethod")?;
51                }
52                write!(f, "{indent}def {}(cls", self.name)?;
53                needs_comma = true;
54            }
55            MethodType::Instance => {
56                write!(f, "{indent}def {}(self", self.name)?;
57                needs_comma = true;
58            }
59        }
60        for arg in &self.args {
61            if needs_comma {
62                write!(f, ", ")?;
63            }
64            write!(f, "{}", arg)?;
65            needs_comma = true;
66        }
67        write!(f, ") -> {}:", self.r#return)?;
68
69        let doc = self.doc;
70        if !doc.is_empty() {
71            writeln!(f)?;
72            let double_indent = format!("{indent}{indent}");
73            docstring::write_docstring(f, self.doc, &double_indent)?;
74        } else {
75            writeln!(f, " ...")?;
76        }
77        Ok(())
78    }
79}