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        write!(f, ") -> {}:", self.r#return)?;
44
45        let doc = self.doc;
46        if !doc.is_empty() {
47            writeln!(f)?;
48            docstring::write_docstring(f, self.doc, indent())?;
49        } else {
50            writeln!(f, " ...")?;
51        }
52        writeln!(f)?;
53        Ok(())
54    }
55}