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 parameters: Parameters,
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    pub is_overload: bool,
17    /// Source file location for deterministic ordering
18    pub file: &'static str,
19    pub line: u32,
20    pub column: u32,
21    /// Index for ordering multiple functions from the same macro invocation
22    pub index: usize,
23}
24
25impl Import for FunctionDef {
26    fn import(&self) -> HashSet<ImportRef> {
27        let mut import = self.r#return.import.clone();
28        import.extend(self.parameters.import());
29        // Add typing_extensions import if deprecated
30        if self.deprecated.is_some() {
31            import.insert("typing_extensions".into());
32        }
33        import
34    }
35}
36
37impl From<&PyFunctionInfo> for FunctionDef {
38    fn from(info: &PyFunctionInfo) -> Self {
39        Self {
40            name: info.name,
41            parameters: Parameters::from_infos(info.parameters),
42            r#return: (info.r#return)(),
43            doc: info.doc,
44            is_async: info.is_async,
45            deprecated: info.deprecated.clone(),
46            type_ignored: info.type_ignored,
47            is_overload: info.is_overload,
48            file: info.file,
49            line: info.line,
50            column: info.column,
51            index: info.index,
52        }
53    }
54}
55
56impl fmt::Display for FunctionDef {
57    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
58        // Add deprecated decorator if present
59        if let Some(deprecated) = &self.deprecated {
60            writeln!(f, "{deprecated}")?;
61        }
62
63        let async_ = if self.is_async { "async " } else { "" };
64        write!(
65            f,
66            "{async_}def {}({}) -> {}:",
67            self.name, self.parameters, self.r#return
68        )?;
69
70        // Calculate type: ignore comment once
71        let type_ignore_comment = if let Some(target) = &self.type_ignored {
72            match target {
73                IgnoreTarget::All => Some("  # type: ignore".to_string()),
74                IgnoreTarget::Specified(rules) => {
75                    let rules_str = rules
76                        .iter()
77                        .map(|r| {
78                            let result = r.parse::<RuleName>().unwrap();
79                            if let RuleName::Custom(custom) = &result {
80                                log::warn!("Unknown custom rule name '{custom}' used in type ignore. Ensure this is intended.");
81                            }
82                            result
83                        })
84                        .join(",");
85                    Some(format!("  # type: ignore[{rules_str}]"))
86                }
87            }
88        } else {
89            None
90        };
91
92        let doc = self.doc;
93        if !doc.is_empty() {
94            // Add type: ignore comment for functions with docstrings
95            if let Some(comment) = &type_ignore_comment {
96                write!(f, "{comment}")?;
97            }
98            writeln!(f)?;
99            docstring::write_docstring(f, self.doc, indent())?;
100        } else {
101            write!(f, " ...")?;
102            // Add type: ignore comment for functions without docstrings
103            if let Some(comment) = &type_ignore_comment {
104                write!(f, "{comment}")?;
105            }
106            writeln!(f)?;
107        }
108        writeln!(f)?;
109        Ok(())
110    }
111}
pyo3_stub_gen/generate/function.rs

pyo3_stub_gen/generate/
function.rs
