pyo3_stub_gen/generate/
method.rs1use 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#[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 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 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 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 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 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 if let Some(comment) = &type_ignore_comment {
122 write!(f, "{}", comment)?;
123 }
124 writeln!(f)?;
125 }
126 Ok(())
127 }
128}