python_ast/ast/tree/
compare.rs

1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, Bound, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods, types::PyTypeMethods};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{
7    dump, extraction_failure, err_from, CodeGen, CodeGenContext, CompareNotYetImplemented, ExprType,
8    PythonOptions, SymbolTableScopes,
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12pub enum Compares {
13    Eq,
14    NotEq,
15    Lt,
16    LtE,
17    Gt,
18    GtE,
19    Is,
20    IsNot,
21    In,
22    NotIn,
23
24    Unknown,
25}
26
27#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
28pub struct Compare {
29    pub ops: Vec<Compares>,
30    pub left: Box<ExprType>,
31    pub comparators: Vec<ExprType>,
32}
33
34impl<'a, 'py> FromPyObject<'a, 'py> for Compare {
35    type Error = pyo3::PyErr;
36    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
37        tracing::debug!("ob: {}", dump(&ob, None)?);
38
39        // Python allows for multiple comparators, rust we only supports one, so we have to rewrite the comparison a little.
40        let ops_bound: Vec<Bound<PyAny>> = ob
41            .getattr("ops")
42            .map_err(|e| extraction_failure("comparison operators", &ob, e))?
43            .extract()
44            .map_err(|e| extraction_failure("comparison operators", &ob, e))?;
45
46        let mut op_list = Vec::new();
47
48        for op in ops_bound.iter() {
49            let op_type = op
50                .get_type()
51                .name()
52                .map_err(|e| extraction_failure("comparison operator type", &ob, e))?;
53
54            let op_type_str: String = op_type.extract()?;
55            let op = match op_type_str.as_str() {
56                "Eq" => Compares::Eq,
57                "NotEq" => Compares::NotEq,
58                "Lt" => Compares::Lt,
59                "LtE" => Compares::LtE,
60                "Gt" => Compares::Gt,
61                "GtE" => Compares::GtE,
62                "Is" => Compares::Is,
63                "IsNot" => Compares::IsNot,
64                "In" => Compares::In,
65                "NotIn" => Compares::NotIn,
66
67                _ => {
68                    tracing::debug!("Found unknown Compare with type: {}", op_type_str);
69                    Compares::Unknown
70                }
71            };
72            op_list.push(op);
73        }
74
75        let left = ob.getattr("left").map_err(|e| extraction_failure("left", &ob, e))?;
76
77        let comparators = ob.getattr("comparators").map_err(|e| extraction_failure("comparators", &ob, e))?;
78        tracing::debug!(
79            "left: {}, comparators: {}",
80            dump(&left, None)?,
81            dump(&comparators, None)?
82        );
83
84        let left = left.extract().map_err(|e| extraction_failure("getting binary operator operand", &ob, e))?;
85        let comparators: Vec<ExprType> = comparators
86            .extract()
87            .map_err(|e| extraction_failure("comparators", &ob, e))?;
88
89        tracing::debug!(
90            "left: {:?}, comparators: {:?}, op: {:?}",
91            left,
92            comparators,
93            op_list
94        );
95
96        return Ok(Compare {
97            ops: op_list,
98            left: Box::new(left),
99            comparators: comparators,
100        });
101    }
102}
103
104impl CodeGen for Compare {
105    type Context = CodeGenContext;
106    type Options = PythonOptions;
107    type SymbolTable = SymbolTableScopes;
108
109    fn to_rust(
110        self,
111        ctx: Self::Context,
112        options: Self::Options,
113        symbols: Self::SymbolTable,
114    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
115        // A CHAINED comparison (`a < b < c`) must evaluate every operand
116        // exactly once — the naive `a < b && b < c` expansion evaluates
117        // `b` twice, running its side effects twice and, for a
118        // non-deterministic operand, even yielding a different answer
119        // than Python. Bind each operand to a temporary at the point
120        // Python evaluates it, and nest the remaining tests inside the
121        // `&&` so a false prefix leaves later operands unevaluated, as
122        // Python's short circuit does. The temporaries bind by
123        // REFERENCE so an operand that is a live variable is not moved
124        // out of the enclosing scope.
125        if self.ops.len() > 1 {
126            return self.to_rust_chained(ctx, options, symbols);
127        }
128        let mut outer_ts = TokenStream::new();
129        // Python chains comparisons pairwise: `a < b < c` means
130        // `a < b && b < c`, so each comparator becomes the left operand of
131        // the next comparison.
132        let mut left = self
133            .left
134            .clone()
135            .to_rust(ctx.clone(), options.clone(), symbols.clone())?;
136        let ops = self.ops.clone();
137        let comparators = self.comparators.clone();
138
139        let mut index = 0;
140        for op in ops.iter() {
141            let comparator_ast = comparators
142                .get(index)
143                .ok_or("comparison has more operators than comparators")?;
144            // The operand AST feeding this comparison's left side: the
145            // original left for the first op, the previous comparator after.
146            let left_ast = if index == 0 {
147                self.left.as_ref()
148            } else {
149                &comparators[index - 1]
150            };
151            // `x is None` / `x is not None` test None-ness, not equality:
152            // Option values report is_none(), plain values are never None.
153            if matches!(op, Compares::Is | Compares::IsNot) {
154                let none_check = if crate::is_none_expr(comparator_ast) {
155                    Some(left_ast)
156                } else if crate::is_none_expr(left_ast) {
157                    Some(comparator_ast)
158                } else {
159                    None
160                };
161                if let Some(operand) = none_check {
162                    let operand_tokens = operand
163                        .clone()
164                        .to_rust(ctx.clone(), options.clone(), symbols.clone())?;
165                    let tokens = match op {
166                        Compares::Is => quote!((#operand_tokens).py_is_none()),
167                        _ => quote!(!(#operand_tokens).py_is_none()),
168                    };
169                    index += 1;
170                    left = quote!(#operand_tokens);
171                    outer_ts.extend(tokens);
172                    if index < ops.len() {
173                        outer_ts.extend(quote!( && ));
174                    }
175                    continue;
176                }
177            }
178            let comparator = comparator_ast
179                .clone()
180                .to_rust(ctx.clone(), options.clone(), symbols.clone())?;
181            let tokens = match op {
182                Compares::Eq => quote!((#left) == (#comparator)),
183                Compares::NotEq => quote!((#left) != (#comparator)),
184                Compares::Lt => quote!((#left) < (#comparator)),
185                Compares::LtE => quote!((#left) <= (#comparator)),
186                Compares::Gt => quote!((#left) > (#comparator)),
187                Compares::GtE => quote!((#left) >= (#comparator)),
188                Compares::Is => quote!(&#left == &#comparator),
189                Compares::IsNot => quote!(&#left != &#comparator),
190                // Python `in` dispatches on the container: substring for
191                // strings, key lookup for dicts, element lookup for
192                // sequences. The stdpython PyContains trait models that.
193                Compares::In => quote!((#comparator).py_contains(&(#left))),
194                Compares::NotIn => quote!(!(#comparator).py_contains(&(#left))),
195
196                _ => return Err(err_from(CompareNotYetImplemented(self)).into()),
197            };
198
199            index += 1;
200            left = comparator;
201
202            outer_ts.extend(tokens);
203            if index < ops.len() {
204                outer_ts.extend(quote!( && ));
205            }
206        }
207        Ok(outer_ts)
208    }
209}
210
211impl Compare {
212    /// Lower `a OP b OP c ...` with each operand evaluated exactly once
213    /// and Python's short-circuit order preserved:
214    ///
215    /// ```text
216    /// { let t0 = &a; let t1 = &b; t0 OP t1 && { let t2 = &c; t1 OP t2 } }
217    /// ```
218    fn to_rust_chained(
219        self,
220        ctx: CodeGenContext,
221        options: PythonOptions,
222        symbols: SymbolTableScopes,
223    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
224        let mut operands: Vec<&ExprType> = Vec::with_capacity(self.comparators.len() + 1);
225        operands.push(self.left.as_ref());
226        operands.extend(self.comparators.iter());
227
228        let mut rendered = Vec::with_capacity(operands.len());
229        for operand in &operands {
230            rendered.push((*operand).clone().to_rust(
231                ctx.clone(),
232                options.clone(),
233                symbols.clone(),
234            )?);
235        }
236        let names: Vec<proc_macro2::Ident> = (0..operands.len())
237            .map(|i| quote::format_ident!("__rython_cmp{}", i))
238            .collect();
239
240        // A None literal is side-effect free and has no nameable type of
241        // its own, so it is never bound to a temporary; `is None` tests
242        // consume only the other side.
243        let is_none: Vec<bool> = operands.iter().map(|e| crate::is_none_expr(e)).collect();
244        let bind = |i: usize| -> TokenStream {
245            if is_none[i] {
246                return quote!();
247            }
248            let name = &names[i];
249            let value = &rendered[i];
250            quote!(let #name = &(#value);)
251        };
252
253        // The comparison for one link of the chain, over the temporaries.
254        let compare_pair = |i: usize| -> Result<TokenStream, Box<dyn std::error::Error>> {
255            let op = &self.ops[i];
256            let (l, r) = (&names[i], &names[i + 1]);
257            if matches!(op, Compares::Is | Compares::IsNot) {
258                let operand = if is_none[i + 1] {
259                    Some(l)
260                } else if is_none[i] {
261                    Some(r)
262                } else {
263                    None
264                };
265                if let Some(operand) = operand {
266                    return Ok(match op {
267                        Compares::Is => quote!((#operand).py_is_none()),
268                        _ => quote!(!(#operand).py_is_none()),
269                    });
270                }
271            }
272            Ok(match op {
273                Compares::Eq => quote!((#l) == (#r)),
274                Compares::NotEq => quote!((#l) != (#r)),
275                Compares::Lt => quote!((#l) < (#r)),
276                Compares::LtE => quote!((#l) <= (#r)),
277                Compares::Gt => quote!((#l) > (#r)),
278                Compares::GtE => quote!((#l) >= (#r)),
279                Compares::Is => quote!((#l) == (#r)),
280                Compares::IsNot => quote!((#l) != (#r)),
281                Compares::In => quote!((#r).py_contains(#l)),
282                Compares::NotIn => quote!(!(#r).py_contains(#l)),
283                _ => return Err(err_from(CompareNotYetImplemented(self.clone())).into()),
284            })
285        };
286
287        // Build inside out so each operand is bound immediately before
288        // the test that first needs it.
289        let mut acc: Option<TokenStream> = None;
290        for i in (0..self.ops.len()).rev() {
291            let rhs_bind = bind(i + 1);
292            let test = compare_pair(i)?;
293            acc = Some(match acc {
294                None => quote!({ #rhs_bind #test }),
295                Some(rest) => quote!({ #rhs_bind #test && #rest }),
296            });
297        }
298        let first_bind = bind(0);
299        let body = acc.expect("a chained comparison has at least one operator");
300        Ok(quote!({ #first_bind #body }))
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307
308    #[test]
309    fn test_simple_eq() {
310        let options = PythonOptions::default();
311        let result = crate::parse("1 == 2", "test_case.py").unwrap();
312        tracing::info!("Python tree: {:?}", result);
313        //info!("{}", result);
314
315        let code = result.to_rust(
316            CodeGenContext::Module("test_case".to_string()),
317            options,
318            SymbolTableScopes::new(),
319        );
320        tracing::info!("module: {:?}", code);
321    }
322
323    #[test]
324    fn test_complex_compare() {
325        let options = PythonOptions::default();
326        let result = crate::parse("1 < a > 6", "test_case.py").unwrap();
327        tracing::info!("Python tree: {:?}", result);
328        //info!("{}", result);
329
330        let code = result.to_rust(
331            CodeGenContext::Module("test_case".to_string()),
332            options,
333            SymbolTableScopes::new(),
334        );
335        tracing::info!("module: {:?}", code);
336    }
337}
python_ast/ast/tree/compare.rs

python_ast/ast/tree/
compare.rs
