pyo3_stub_gen/generate/
function.rs

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

pyo3_stub_gen/generate/
function.rs
