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    pub is_async: bool,
12    pub deprecated: Option<DeprecatedInfo>,
13}
14
15impl Import for FunctionDef {
16    fn import(&self) -> HashSet<ModuleRef> {
17        let mut import = self.r#return.import.clone();
18        for arg in &self.args {
19            import.extend(arg.import().into_iter());
20        }
21        // Add typing_extensions import if deprecated
22        if self.deprecated.is_some() {
23            import.insert("typing_extensions".into());
24        }
25        import
26    }
27}
28
29impl From<&PyFunctionInfo> for FunctionDef {
30    fn from(info: &PyFunctionInfo) -> Self {
31        Self {
32            name: info.name,
33            args: info.args.iter().map(Arg::from).collect(),
34            r#return: (info.r#return)(),
35            doc: info.doc,
36            is_async: info.is_async,
37            deprecated: info.deprecated.clone(),
38        }
39    }
40}
41
42impl fmt::Display for FunctionDef {
43    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
44        // Add deprecated decorator if present
45        if let Some(deprecated) = &self.deprecated {
46            writeln!(f, "{deprecated}")?;
47        }
48
49        let async_ = if self.is_async { "async " } else { "" };
50        write!(f, "{async_}def {}(", self.name)?;
51        for (i, arg) in self.args.iter().enumerate() {
52            write!(f, "{arg}")?;
53            if i != self.args.len() - 1 {
54                write!(f, ", ")?;
55            }
56        }
57        write!(f, ") -> {}:", self.r#return)?;
58
59        let doc = self.doc;
60        if !doc.is_empty() {
61            writeln!(f)?;
62            docstring::write_docstring(f, self.doc, indent())?;
63        } else {
64            writeln!(f, " ...")?;
65        }
66        writeln!(f)?;
67        Ok(())
68    }
69}