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 args: Vec<Arg>,
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        for arg in &self.args {
25            import.extend(arg.import().into_iter());
26        }
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            args: info.args.iter().map(Arg::from).collect(),
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        }
47    }
48}
49
50impl fmt::Display for MethodDef {
51    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
52        let indent = indent();
53        let mut needs_comma = false;
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        match self.r#type {
62            MethodType::Static => {
63                writeln!(f, "{indent}@staticmethod")?;
64                write!(f, "{indent}{async_}def {}(", self.name)?;
65            }
66            MethodType::Class | MethodType::New => {
67                if self.r#type == MethodType::Class {
68                    // new is a classmethod without the decorator
69                    writeln!(f, "{indent}@classmethod")?;
70                }
71                write!(f, "{indent}{async_}def {}(cls", self.name)?;
72                needs_comma = true;
73            }
74            MethodType::Instance => {
75                write!(f, "{indent}{async_}def {}(self", self.name)?;
76                needs_comma = true;
77            }
78        }
79        for arg in &self.args {
80            if needs_comma {
81                write!(f, ", ")?;
82            }
83            write!(f, "{arg}")?;
84            needs_comma = true;
85        }
86        write!(f, ") -> {}:", self.r#return)?;
87
88        // Calculate type: ignore comment once
89        let type_ignore_comment = if let Some(target) = &self.type_ignored {
90            match target {
91                IgnoreTarget::All => Some("  # type: ignore".to_string()),
92                IgnoreTarget::Specified(rules) => {
93                    let rules_str = rules
94                        .iter()
95                        .map(|r| {
96                            let result = r.parse::<RuleName>().unwrap();
97                            if let RuleName::Custom(custom) = &result {
98                                log::warn!("Unknown custom rule name '{custom}' used in type ignore. Ensure this is intended.");
99                            }
100                            result
101                        })
102                        .join(",");
103                    Some(format!("  # type: ignore[{rules_str}]"))
104                }
105            }
106        } else {
107            None
108        };
109
110        let doc = self.doc;
111        if !doc.is_empty() {
112            // Add type: ignore comment for methods with docstrings
113            if let Some(comment) = &type_ignore_comment {
114                write!(f, "{comment}")?;
115            }
116            writeln!(f)?;
117            let double_indent = format!("{indent}{indent}");
118            docstring::write_docstring(f, self.doc, &double_indent)?;
119        } else {
120            write!(f, " ...")?;
121            // Add type: ignore comment for methods without docstrings
122            if let Some(comment) = &type_ignore_comment {
123                write!(f, "{comment}")?;
124            }
125            writeln!(f)?;
126        }
127        Ok(())
128    }
129}
pyo3_stub_gen/generate/method.rs

pyo3_stub_gen/generate/
method.rs
