Skip to main content

python_ast/ast/tree/
bool_ops.rs

1use 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    /// All operands: Python collapses `a and b and c` into one BoolOp node
44    /// with three values.
45    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        // Python's boolean operators return operands, not booleans; for now we
94        // approximate with Rust's short-circuiting operators, folding every
95        // operand (a BoolOp node can carry more than two).
96        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                // `a or None` yields the Option-model None when `a` is
104                // falsy — dropping the None silently returned the falsy
105                // value instead (`0 or None` must be None, not 0).
106                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        //tracing::info!("{}", result.to_rust().unwrap());
134
135        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        //tracing::info!("{}", result);
151
152        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}