python_ast/ast/tree/
bool_ops.rs1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods, types::PyTypeMethods};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{
7 dump, extraction_failure, err_from, BoolOpNotYetImplemented, CodeGen, CodeGenContext, ExprType,
8 PythonOptions, SymbolTableScopes,
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12pub enum BoolOps {
13 And,
14 Or,
15 Unknown,
16}
17
18impl<'a, 'py> FromPyObject<'a, 'py> for BoolOps {
19 type Error = pyo3::PyErr;
20 fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
21 let op_type = ob
22 .get_type()
23 .name()
24 .map_err(|e| extraction_failure("boolean operator type", &ob, e))?;
25
26 let op_type_str: String = op_type.extract()?;
27 let op = match op_type_str.as_str() {
28 "And" => BoolOps::And,
29 "Or" => BoolOps::Or,
30 _ => {
31 tracing::debug!("Found unknown BoolOp {:?}", op_type_str);
32 BoolOps::Unknown
33 }
34 };
35
36 Ok(op)
37 }
38}
39
40#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
41pub struct BoolOp {
42 pub op: BoolOps,
43 pub values: Vec<ExprType>,
46}
47
48impl<'a, 'py> FromPyObject<'a, 'py> for BoolOp {
49 type Error = pyo3::PyErr;
50 fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
51 tracing::debug!("ob: {}", dump(&ob, None)?);
52 let op = ob.getattr("op").map_err(|e| extraction_failure("op", &ob, e))?;
53
54 let op_type = op
55 .get_type()
56 .name()
57 .map_err(|e| extraction_failure("boolean operator type", &ob, e))?;
58
59 let values = ob.getattr("values").map_err(|e| extraction_failure("values", &ob, e))?;
60
61 tracing::debug!("BoolOps values: {}", dump(&values, None)?);
62
63 let values: Vec<ExprType> = values.extract().map_err(|e| extraction_failure("getting values from BoolOp", &ob, e))?;
64
65 let op_type_str: String = op_type.extract()?;
66 let op = match op_type_str.as_str() {
67 "And" => BoolOps::And,
68 "Or" => BoolOps::Or,
69
70 _ => {
71 tracing::debug!("Found unknown BoolOp {:?}", op);
72 BoolOps::Unknown
73 }
74 };
75
76 tracing::debug!("values: {:?}, op: {:?}/{:?}", values, op_type, op);
77
78 return Ok(BoolOp { op, values });
79 }
80}
81
82impl<'a> CodeGen for BoolOp {
83 type Context = CodeGenContext;
84 type Options = PythonOptions;
85 type SymbolTable = SymbolTableScopes;
86
87 fn to_rust(
88 self,
89 ctx: Self::Context,
90 options: Self::Options,
91 symbols: Self::SymbolTable,
92 ) -> Result<TokenStream, Box<dyn std::error::Error>> {
93 let mut rendered = Vec::new();
97 for value in self.values.clone() {
98 rendered.push(value.to_rust(ctx.clone(), options.clone(), symbols.clone())?);
99 }
100
101 match self.op {
102 BoolOps::Or => {
103 if let Some(last) = rendered.last() {
107 if last.to_string().trim() == "None" && rendered.len() == 2 {
108 let first = &rendered[0];
109 return Ok(quote!({
110 let __rython_or = #first;
111 if (__rython_or).is_truthy() { Some(__rython_or) } else { None }
112 }));
113 }
114 }
115 Ok(quote!(#((#rendered))||*))
116 }
117 BoolOps::And => Ok(quote!(#((#rendered))&&*)),
118
119 _ => Err(err_from(BoolOpNotYetImplemented(self)).into()),
120 }
121 }
122}
123
124#[cfg(test)]
125mod tests {
126 use super::*;
127
128 #[test]
129 fn test_and() {
130 let options = PythonOptions::default();
131 let result = crate::parse("1 and 2", "test_case.py").unwrap();
132 tracing::info!("Python tree: {:?}", result);
133 let code = result
136 .to_rust(
137 CodeGenContext::Module("test_case".to_string()),
138 options,
139 SymbolTableScopes::new(),
140 )
141 .unwrap();
142 tracing::info!("module: {:?}", code);
143 }
144
145 #[test]
146 fn test_or() {
147 let options = PythonOptions::default();
148 let result = crate::parse("1 or 2", "test_case.py").unwrap();
149 tracing::info!("Python tree: {:?}", result);
150 let code = result
153 .to_rust(
154 CodeGenContext::Module("test_case".to_string()),
155 options,
156 SymbolTableScopes::new(),
157 )
158 .unwrap();
159 tracing::info!("module: {:?}", code);
160 }
161}