Skip to main content

python_ast/ast/tree/
named_expression.rs

1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, PyAny, PyResult, FromPyObject, prelude::PyAnyMethods};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{CodeGen, CodeGenContext, ExprType, PythonOptions, SymbolTableScopes};
7
8/// A keyword argument, gnerally used in function calls.
9#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)]
10pub struct NamedExpr {
11    pub left: Box<ExprType>,
12    pub right: Box<ExprType>,
13}
14
15impl<'a, 'py> FromPyObject<'a, 'py> for NamedExpr {
16    type Error = pyo3::PyErr;
17    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
18        // ast.NamedExpr stores its operands as `target` and `value`.
19        let left = ob.getattr("target")?.extract::<ExprType>()?;
20        let right = ob.getattr("value")?.extract::<ExprType>()?;
21        Ok(NamedExpr {
22            left: Box::new(left),
23            right: Box::new(right),
24        })
25    }
26}
27
28impl CodeGen for NamedExpr {
29    type Context = CodeGenContext;
30    type Options = PythonOptions;
31    type SymbolTable = SymbolTableScopes;
32
33    fn to_rust(
34        self,
35        ctx: Self::Context,
36        options: Self::Options,
37        symbols: Self::SymbolTable,
38    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
39        let left = self
40            .left
41            .clone()
42            .to_rust(ctx.clone(), options.clone(), symbols.clone())?;
43        let right = self.right.clone().to_rust(ctx, options, symbols)?;
44        Ok(quote!(#left = #right))
45    }
46}
47
48#[cfg(test)]
49mod test {
50    use super::*;
51    use crate::{Constant, ExprType, Name};
52    use litrs::*;
53
54    #[test]
55    fn test_named_expression() {
56        let named_expression = NamedExpr {
57            left: Box::new(ExprType::Name(Name {
58                id: "a".to_string(),
59            })),
60            right: Box::new(ExprType::Constant(Constant(Some(Literal::Integer(
61                IntegerLit::parse("1".to_string()).unwrap(),
62            ))))),
63        };
64        let rust = named_expression
65            .to_rust(
66                CodeGenContext::Module("test".to_string()),
67                PythonOptions::default(),
68                SymbolTableScopes::new(),
69            )
70            .unwrap();
71        assert_eq!(rust.to_string(), "a = 1");
72    }
73}