python_ast/ast/tree/
function_def.rs

1use log::debug;
2use proc_macro2::TokenStream;
3use pyo3::FromPyObject;
4use quote::{format_ident, quote};
5use serde::{Deserialize, Serialize};
6
7use crate::{
8    CodeGen, CodeGenContext, ExprType, Object, ParameterList, PythonOptions, Statement,
9    StatementType, SymbolTableNode, SymbolTableScopes,
10};
11
12#[derive(Clone, Debug, FromPyObject, Serialize, Deserialize, PartialEq)]
13pub struct FunctionDef {
14    pub name: String,
15    pub args: ParameterList,
16    pub body: Vec<Statement>,
17    pub decorator_list: Vec<String>,
18}
19
20impl CodeGen for FunctionDef {
21    type Context = CodeGenContext;
22    type Options = PythonOptions;
23    type SymbolTable = SymbolTableScopes;
24
25    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
26        let mut symbols = symbols;
27        symbols.insert(
28            self.name.clone(),
29            SymbolTableNode::FunctionDef(self.clone()),
30        );
31        symbols
32    }
33
34    fn to_rust(
35        self,
36        ctx: Self::Context,
37        options: Self::Options,
38        symbols: SymbolTableScopes,
39    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
40        let mut streams = TokenStream::new();
41        let fn_name = format_ident!("{}", self.name);
42
43        // The Python convention is that functions that begin with a single underscore,
44        // it's private. Otherwise, it's public. We formalize that by default.
45        let visibility = if self.name.starts_with("_") && !self.name.starts_with("__") {
46            format_ident!("")
47        } else if self.name.starts_with("__") && self.name.ends_with("__") {
48            format_ident!("pub(crate)")
49        } else {
50            format_ident!("pub")
51        };
52
53        let is_async = match ctx.clone() {
54            CodeGenContext::Async(_) => {
55                quote!(async)
56            }
57            _ => quote!(),
58        };
59
60        let parameters = self
61            .args
62            .clone()
63            .to_rust(ctx.clone(), options.clone(), symbols.clone())
64            .expect(format!("parsing arguments {:?}", self.args).as_str());
65
66        for s in self.body.iter() {
67            streams.extend(
68                s.clone()
69                    .to_rust(ctx.clone(), options.clone(), symbols.clone())
70                    .expect(format!("parsing statement {:?}", s).as_str()),
71            );
72            streams.extend(quote!(;));
73        }
74
75        let docstring = if let Some(d) = self.get_docstring() {
76            format!("{}", d)
77        } else {
78            "".to_string()
79        };
80
81        let function = quote! {
82            #[doc = #docstring]
83            #visibility #is_async fn #fn_name(#parameters) {
84                #streams
85            }
86        };
87
88        debug!("function: {}", function);
89        Ok(function)
90    }
91
92    fn get_docstring(&self) -> Option<String> {
93        let expr = self.body[0].clone();
94        match expr.statement {
95            StatementType::Expr(e) => match e.value {
96                ExprType::Constant(c) => Some(c.to_string()),
97                _ => None,
98            },
99            _ => None,
100        }
101    }
102}
103
104impl Object for FunctionDef {}
python_ast/ast/tree/function_def.rs

python_ast/ast/tree/
function_def.rs
