pyo3_stub_gen/generate/
function.rs

1use crate::{generate::*, type_info::*, TypeInfo};
2use std::fmt;
3
4/// Definition of a Python function.
5#[derive(Debug, Clone, PartialEq)]
6pub struct FunctionDef {
7    pub name: &'static str,
8    pub args: Vec<Arg>,
9    pub r#return: TypeInfo,
10    pub doc: &'static str,
11}
12
13impl Import for FunctionDef {
14    fn import(&self) -> HashSet<ModuleRef> {
15        let mut import = self.r#return.import.clone();
16        for arg in &self.args {
17            import.extend(arg.import().into_iter());
18        }
19        import
20    }
21}
22
23impl From<&PyFunctionInfo> for FunctionDef {
24    fn from(info: &PyFunctionInfo) -> Self {
25        Self {
26            name: info.name,
27            args: info.args.iter().map(Arg::from).collect(),
28            r#return: (info.r#return)(),
29            doc: info.doc,
30        }
31    }
32}
33
34impl fmt::Display for FunctionDef {
35    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
36        write!(f, "def {}(", self.name)?;
37        for (i, arg) in self.args.iter().enumerate() {
38            write!(f, "{}", arg)?;
39            if i != self.args.len() - 1 {
40                write!(f, ", ")?;
41            }
42        }
43        writeln!(f, ") -> {}:", self.r#return)?;
44
45        let doc = self.doc;
46        let indent = indent();
47        if !doc.is_empty() {
48            writeln!(f, r#"{indent}r""""#)?;
49            for line in doc.lines() {
50                writeln!(f, "{indent}{}", line)?;
51            }
52            writeln!(f, r#"{indent}""""#)?;
53        }
54        writeln!(f, "{indent}...")?;
55        writeln!(f)?;
56        Ok(())
57    }
58}