pyo3_stub_gen/generate/
method.rs

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