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