pyo3_stub_gen/generate/
method.rs

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