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

pyo3_stub_gen/generate/
function.rs
