pyo3_stub_gen/generate/
function.rs

1use crate::stub_type::ImportRef;
2use crate::{generate::*, rule_name::RuleName, type_info::*, TypeInfo};
3use itertools::Itertools;
4use std::fmt;
5
6/// Definition of a Python function.
7#[derive(Debug, Clone, PartialEq)]
8pub struct FunctionDef {
9    pub name: &'static str,
10    pub args: Vec<Arg>,
11    pub r#return: TypeInfo,
12    pub doc: &'static str,
13    pub is_async: bool,
14    pub deprecated: Option<DeprecatedInfo>,
15    pub type_ignored: Option<IgnoreTarget>,
16}
17
18impl Import for FunctionDef {
19    fn import(&self) -> HashSet<ImportRef> {
20        let mut import = self.r#return.import.clone();
21        for arg in &self.args {
22            import.extend(arg.import().into_iter());
23        }
24        // Add typing_extensions import if deprecated
25        if self.deprecated.is_some() {
26            import.insert("typing_extensions".into());
27        }
28        import
29    }
30}
31
32impl From<&PyFunctionInfo> for FunctionDef {
33    fn from(info: &PyFunctionInfo) -> Self {
34        Self {
35            name: info.name,
36            args: info.args.iter().map(Arg::from).collect(),
37            r#return: (info.r#return)(),
38            doc: info.doc,
39            is_async: info.is_async,
40            deprecated: info.deprecated.clone(),
41            type_ignored: info.type_ignored,
42        }
43    }
44}
45
46impl fmt::Display for FunctionDef {
47    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
48        // Add deprecated decorator if present
49        if let Some(deprecated) = &self.deprecated {
50            writeln!(f, "{deprecated}")?;
51        }
52
53        let async_ = if self.is_async { "async " } else { "" };
54        write!(f, "{async_}def {}(", self.name)?;
55        for (i, arg) in self.args.iter().enumerate() {
56            write!(f, "{arg}")?;
57            if i != self.args.len() - 1 {
58                write!(f, ", ")?;
59            }
60        }
61        write!(f, ") -> {}:", self.r#return)?;
62
63        // Calculate type: ignore comment once
64        let type_ignore_comment = if let Some(target) = &self.type_ignored {
65            match target {
66                IgnoreTarget::All => Some("  # type: ignore".to_string()),
67                IgnoreTarget::Specified(rules) => {
68                    let rules_str = rules
69                        .iter()
70                        .map(|r| {
71                            let result = r.parse::<RuleName>().unwrap();
72                            if let RuleName::Custom(custom) = &result {
73                                log::warn!("Unknown custom rule name '{custom}' used in type ignore. Ensure this is intended.");
74                            }
75                            result
76                        })
77                        .join(",");
78                    Some(format!("  # type: ignore[{rules_str}]"))
79                }
80            }
81        } else {
82            None
83        };
84
85        let doc = self.doc;
86        if !doc.is_empty() {
87            // Add type: ignore comment for functions with docstrings
88            if let Some(comment) = &type_ignore_comment {
89                write!(f, "{comment}")?;
90            }
91            writeln!(f)?;
92            docstring::write_docstring(f, self.doc, indent())?;
93        } else {
94            write!(f, " ...")?;
95            // Add type: ignore comment for functions without docstrings
96            if let Some(comment) = &type_ignore_comment {
97                write!(f, "{comment}")?;
98            }
99            writeln!(f)?;
100        }
101        writeln!(f)?;
102        Ok(())
103    }
104}
pyo3_stub_gen/generate/function.rs

pyo3_stub_gen/generate/
function.rs
