python_ast/ast/tree/
if_exp.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, PyAttributeExtractor
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12pub struct IfExp {
13 pub test: Box<ExprType>,
14 pub body: Box<ExprType>,
15 pub orelse: Box<ExprType>,
16 pub lineno: Option<usize>,
17 pub col_offset: Option<usize>,
18 pub end_lineno: Option<usize>,
19 pub end_col_offset: Option<usize>,
20}
21
22impl<'a, 'py> FromPyObject<'a, 'py> for IfExp {
23 type Error = pyo3::PyErr;
24 fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
25 let test = ob.extract_attr_with_context("test", "if expression test")?;
26 let body = ob.extract_attr_with_context("body", "if expression body")?;
27 let orelse = ob.extract_attr_with_context("orelse", "if expression orelse")?;
28
29 let test = test.extract().map_err(|e| extraction_failure("getting if expression test", &ob, e))?;
30 let body = body.extract().map_err(|e| extraction_failure("getting if expression body", &ob, e))?;
31 let orelse = orelse.extract().map_err(|e| extraction_failure("getting if expression orelse", &ob, e))?;
32
33 Ok(IfExp {
34 test: Box::new(test),
35 body: Box::new(body),
36 orelse: Box::new(orelse),
37 lineno: ob.lineno(),
38 col_offset: ob.col_offset(),
39 end_lineno: ob.end_lineno(),
40 end_col_offset: ob.end_col_offset(),
41 })
42 }
43}
44
45impl_node_with_positions!(IfExp { lineno, col_offset, end_lineno, end_col_offset });
46
47impl CodeGen for IfExp {
48 type Context = CodeGenContext;
49 type Options = PythonOptions;
50 type SymbolTable = SymbolTableScopes;
51
52 fn to_rust(
53 self,
54 ctx: Self::Context,
55 options: Self::Options,
56 symbols: Self::SymbolTable,
57 ) -> Result<TokenStream, Box<dyn std::error::Error>> {
58 let test =
59 crate::condition_to_rust(&self.test, ctx.clone(), options.clone(), symbols.clone())?;
60 let body = self.body.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
61 let orelse = self.orelse.to_rust(ctx, options, symbols)?;
62
63 Ok(quote! {
64 if #test { #body } else { #orelse }
65 })
66 }
67}
68
69#[cfg(test)]
70mod tests {
71 use super::*;
72 use crate::create_parse_test;
73
74 create_parse_test!(test_if_expression, "x if condition else y", "if_exp_test.py");
75 create_parse_test!(test_nested_if_expression, "a if b else c if d else e", "if_exp_test.py");
76}