1use crate::{generate::*, type_info::*, TypeInfo};
2use std::fmt;
3
4/// Definition of a Python function.
5#[derive(Debug, Clone, PartialEq)]
6pub struct FunctionDef {
7 pub name: &'static str,
8 pub args: Vec<Arg>,
9 pub r#return: TypeInfo,
10 pub doc: &'static str,
11}
12
13impl Import for FunctionDef {
14 fn import(&self) -> HashSet<ModuleRef> {
15 let mut import: HashSet = self.r#return.import.clone();
16 for arg: &Arg in &self.args {
17 import.extend(arg.import().into_iter());
18 }
19 import
20 }
21}
22
23impl From<&PyFunctionInfo> for FunctionDef {
24 fn from(info: &PyFunctionInfo) -> Self {
25 Self {
26 name: info.name,
27 args: info.args.iter().map(Arg::from).collect(),
28 r#return: (info.r#return)(),
29 doc: info.doc,
30 }
31 }
32}
33
34impl fmt::Display for FunctionDef {
35 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
36 write!(f, "def {}(", self.name)?;
37 for (i, arg) in self.args.iter().enumerate() {
38 write!(f, "{}", arg)?;
39 if i != self.args.len() - 1 {
40 write!(f, ", ")?;
41 }
42 }
43 writeln!(f, ") -> {}:", self.r#return)?;
44
45 let doc = self.doc;
46 let indent = indent();
47 if !doc.is_empty() {
48 writeln!(f, r#"{indent}r""""#)?;
49 for line in doc.lines() {
50 writeln!(f, "{indent}{}", line)?;
51 }
52 writeln!(f, r#"{indent}""""#)?;
53 }
54 writeln!(f, "{indent}...")?;
55 writeln!(f)?;
56 Ok(())
57 }
58}
59