python_ast/ast/tree/
lambda.rs1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, FromPyObject, PyAny, PyResult, types::PyAnyMethods};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{extraction_failure,
7 CodeGen, CodeGenContext, ExprType, PythonOptions, SymbolTableScopes,
8 Node, impl_node_with_positions, ParameterList, PyAttributeExtractor
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12pub struct Lambda {
13 pub args: ParameterList,
14 pub body: Box<ExprType>,
15 pub lineno: Option<usize>,
16 pub col_offset: Option<usize>,
17 pub end_lineno: Option<usize>,
18 pub end_col_offset: Option<usize>,
19}
20
21impl<'a, 'py> FromPyObject<'a, 'py> for Lambda {
22 type Error = pyo3::PyErr;
23 fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
24 let args = ob.extract_attr_with_context("args", "lambda arguments")?;
25 let body = ob.extract_attr_with_context("body", "lambda body")?;
26
27 let args = args.extract().map_err(|e| extraction_failure("getting lambda arguments", &ob, e))?;
28 let body = body.extract().map_err(|e| extraction_failure("getting lambda body", &ob, e))?;
29
30 Ok(Lambda {
31 args,
32 body: Box::new(body),
33 lineno: ob.lineno(),
34 col_offset: ob.col_offset(),
35 end_lineno: ob.end_lineno(),
36 end_col_offset: ob.end_col_offset(),
37 })
38 }
39}
40
41impl_node_with_positions!(Lambda { lineno, col_offset, end_lineno, end_col_offset });
42
43impl CodeGen for Lambda {
44 type Context = CodeGenContext;
45 type Options = PythonOptions;
46 type SymbolTable = SymbolTableScopes;
47
48 fn to_rust(
49 self,
50 ctx: Self::Context,
51 options: Self::Options,
52 symbols: Self::SymbolTable,
53 ) -> Result<TokenStream, Box<dyn std::error::Error>> {
54 let params: Vec<_> = self
58 .args
59 .args
60 .iter()
61 .map(|param| crate::safe_ident(¶m.arg))
62 .collect();
63 let body = self.body.to_rust(ctx, options, symbols)?;
64
65 Ok(quote! {
66 |#(#params),*| #body
67 })
68 }
69}
70
71#[cfg(test)]
72mod tests {
73 use super::*;
74 use crate::create_parse_test;
75
76 create_parse_test!(test_simple_lambda, "lambda x: x + 1", "lambda_test.py");
77 create_parse_test!(test_lambda_with_args, "lambda x, y: x * y", "lambda_test.py");
78 create_parse_test!(test_lambda_no_args, "lambda: 42", "lambda_test.py");
79}