python_ast/ast/tree/
bin_ops.rs

1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{
7    dump, extraction_failure, err_from, BinOpNotYetImplemented, BinaryOperation, CodeGen, CodeGenContext, ExprType,
8    FromPythonString, PyAttributeExtractor, PythonOperator, PythonOptions, SymbolTableScopes,
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12pub enum BinOps {
13    Add,
14    Sub,
15    Mult,
16    Div,
17    FloorDiv,
18    Mod,
19    Pow,
20    LShift,
21    RShift,
22    BitOr,
23    BitXor,
24    BitAnd,
25    MatMult,
26
27    Unknown,
28}
29
30impl FromPythonString for BinOps {
31    fn from_python_string(s: &str) -> Option<Self> {
32        match s {
33            "Add" => Some(BinOps::Add),
34            "Sub" => Some(BinOps::Sub),
35            "Mult" => Some(BinOps::Mult),
36            "Div" => Some(BinOps::Div),
37            "FloorDiv" => Some(BinOps::FloorDiv),
38            "Mod" => Some(BinOps::Mod),
39            "Pow" => Some(BinOps::Pow),
40            "LShift" => Some(BinOps::LShift),
41            "RShift" => Some(BinOps::RShift),
42            "BitOr" => Some(BinOps::BitOr),
43            "BitXor" => Some(BinOps::BitXor),
44            "BitAnd" => Some(BinOps::BitAnd),
45            "MatMult" => Some(BinOps::MatMult),
46            _ => None,
47        }
48    }
49    
50    fn unknown() -> Self {
51        BinOps::Unknown
52    }
53}
54
55impl PythonOperator for BinOps {
56    fn to_rust_op(&self) -> Result<TokenStream, Box<dyn std::error::Error>> {
57        match self {
58            BinOps::Add => Ok(quote!(+)),
59            BinOps::Sub => Ok(quote!(-)),
60            BinOps::Mult => Ok(quote!(*)),
61            BinOps::Div => Ok(quote!(as f64 /)),
62            BinOps::FloorDiv => Ok(quote!(/)),
63            BinOps::Mod => Ok(quote!(%)),
64            BinOps::Pow => Ok(quote!(.pow)),
65            BinOps::LShift => Ok(quote!(<<)),
66            BinOps::RShift => Ok(quote!(>>)),
67            BinOps::BitOr => Ok(quote!(|)),
68            BinOps::BitXor => Ok(quote!(^)),
69            BinOps::BitAnd => Ok(quote!(&)),
70            _ => Err(err_from(BinOpNotYetImplemented(BinOp {
71                op: self.clone(),
72                left: Box::new(ExprType::Name(crate::Name { id: "unknown".to_string() })),
73                right: Box::new(ExprType::Name(crate::Name { id: "unknown".to_string() })),
74            })).into()),
75        }
76    }
77    
78    fn is_unknown(&self) -> bool {
79        matches!(self, BinOps::Unknown)
80    }
81}
82
83#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
84pub struct BinOp {
85    pub op: BinOps,
86    pub left: Box<ExprType>,
87    pub right: Box<ExprType>,
88}
89
90impl BinaryOperation for BinOp {
91    type OperatorType = BinOps;
92    
93    fn operator(&self) -> &Self::OperatorType {
94        &self.op
95    }
96    
97    fn left(&self) -> &ExprType {
98        &self.left
99    }
100    
101    fn right(&self) -> &ExprType {
102        &self.right
103    }
104}
105
106impl<'a, 'py> FromPyObject<'a, 'py> for BinOp {
107    type Error = pyo3::PyErr;
108    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
109        tracing::debug!("ob: {}", dump(&ob, None)?);
110        
111        let op = ob.extract_attr_with_context("op", "binary operator")?;
112        let op_type_str = op.extract_type_name("binary operator")?;
113        
114        let left = ob.extract_attr_with_context("left", "binary operand")?;
115        let right = ob.extract_attr_with_context("right", "binary operand")?;
116        
117        tracing::debug!("left: {}, right: {}", dump(&left, None)?, dump(&right, None)?);
118
119        let op = BinOps::parse_or_unknown(&op_type_str);
120        if matches!(op, BinOps::Unknown) {
121            tracing::debug!("Found unknown BinOp {:?}", op_type_str);
122        }
123
124        let left = left.extract().map_err(|e| extraction_failure("getting binary operator operand", &ob, e))?;
125        let right = right.extract().map_err(|e| extraction_failure("getting binary operator operand", &ob, e))?;
126
127        Ok(BinOp {
128            op,
129            left: Box::new(left),
130            right: Box::new(right),
131        })
132    }
133}
134
135impl CodeGen for BinOp {
136    type Context = CodeGenContext;
137    type Options = PythonOptions;
138    type SymbolTable = SymbolTableScopes;
139
140    fn to_rust(
141        self,
142        ctx: Self::Context,
143        options: Self::Options,
144        symbols: Self::SymbolTable,
145    ) -> std::result::Result<TokenStream, Box<dyn std::error::Error>> {
146        // Python's ** promotes based on operand types; route through the
147        // stdpython py_pow helper, which implements those semantics.
148        if matches!(self.op, BinOps::Pow) {
149            let left = self.left.clone().to_rust(ctx.clone(), options.clone(), symbols.clone())?;
150            let right = self.right.clone().to_rust(ctx, options, symbols)?;
151            return Ok(quote!(py_pow(#left, #right)));
152        }
153        
154        // For Div, we need to cast to f64
155        if matches!(self.op, BinOps::Div) {
156            let left = self.left.clone().to_rust(ctx.clone(), options.clone(), symbols.clone())?;
157            let right = self.right.clone().to_rust(ctx, options, symbols)?;
158            return Ok(quote!((#left) as f64 / (#right) as f64));
159        }
160
161        // Python's // floors toward negative infinity and % takes the
162        // divisor's sign; Rust's / and % truncate. Route through the
163        // stdpython helpers, which implement the Python semantics.
164        if matches!(self.op, BinOps::FloorDiv) {
165            let left = self.left.clone().to_rust(ctx.clone(), options.clone(), symbols.clone())?;
166            let right = self.right.clone().to_rust(ctx, options, symbols)?;
167            return Ok(quote!(py_floordiv(#left, #right)));
168        }
169
170        if matches!(self.op, BinOps::Mod) {
171            let left = self.left.clone().to_rust(ctx.clone(), options.clone(), symbols.clone())?;
172            let right = self.right.clone().to_rust(ctx, options, symbols)?;
173            return Ok(quote!(py_mod(#left, #right)));
174        }
175        
176        // Python's * repeats sequences when one operand is a string:
177        // "!" * 3 == "!!!". Route literal-string repetition through the
178        // stdpython multiply_string helper (numeric multiplication keeps
179        // the plain operator below).
180        if matches!(self.op, BinOps::Mult) {
181            let left_is_str = matches!(&*self.left, ExprType::Constant(c) if matches!(&c.0, Some(litrs::Literal::String(_))))
182                || matches!(&*self.left, ExprType::JoinedStr(_));
183            let right_is_str = matches!(&*self.right, ExprType::Constant(c) if matches!(&c.0, Some(litrs::Literal::String(_))))
184                || matches!(&*self.right, ExprType::JoinedStr(_));
185            if left_is_str || right_is_str {
186                let left = self.left.clone().to_rust(ctx.clone(), options.clone(), symbols.clone())?;
187                let right = self.right.clone().to_rust(ctx, options, symbols)?;
188                return Ok(if left_is_str {
189                    quote!(multiply_string(#left, (#right) as i64))
190                } else {
191                    quote!(multiply_string(#right, (#left) as i64))
192                });
193            }
194        }
195
196        // Python `+` covers cases Rust's Add doesn't (String + String,
197        // int/float promotion, list concatenation): lower through the
198        // stdpython PyAdd trait, which borrows both operands. Bare numeric
199        // literals get an explicit type: trait-method resolution on an
200        // unanchored `{integer}` receiver fails before literal fallback
201        // (e.g. `1 + 2` in an unannotated position).
202        if matches!(self.op, BinOps::Add) {
203            let left = self.left.clone().to_rust(ctx.clone(), options.clone(), symbols.clone())?;
204            let right = self.right.clone().to_rust(ctx, options, symbols)?;
205            let left = anchor_numeric_literal(&self.left, left);
206            let right = anchor_numeric_literal(&self.right, right);
207            return Ok(quote!((#left).py_add(&(#right))));
208        }
209
210        // Use the generic binary operation implementation for everything else
211        self.generate_rust_code(ctx, options, symbols)
212    }
213}
214
215/// Give a bare numeric literal (possibly under unary +/-) a concrete type
216/// so PyAdd's trait-method resolution has an anchored receiver: int
217/// literals become i64, float literals f64. Anything else is left to its
218/// own type.
219fn anchor_numeric_literal(expr: &ExprType, tokens: TokenStream) -> TokenStream {
220    fn literal_type(expr: &ExprType) -> Option<TokenStream> {
221        match expr {
222            ExprType::Constant(c) => match &c.0 {
223                Some(litrs::Literal::Integer(_)) => Some(quote!(i64)),
224                Some(litrs::Literal::Float(_)) => Some(quote!(f64)),
225                _ => None,
226            },
227            ExprType::UnaryOp(u) => literal_type(&u.operand),
228            _ => None,
229        }
230    }
231    match literal_type(expr) {
232        Some(ty) => quote!((#tokens) as #ty),
233        None => tokens,
234    }
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use crate::create_parse_test;
241
242    create_parse_test!(test_add, "1 + 2", "test_case.py");
243    create_parse_test!(test_subtract, "1 - 2", "test_case.py");
244    create_parse_test!(test_multiply, "3 * 4", "test_case.py");
245    create_parse_test!(test_divide, "8 / 2", "test_case.py");
246    create_parse_test!(test_power, "2 ** 3", "test_case.py");
247    create_parse_test!(test_modulo, "10 % 3", "test_case.py");
248    
249    #[test]
250    fn test_unknown_operator() {
251        let unknown_op = BinOps::Unknown;
252        assert!(unknown_op.is_unknown());
253        assert!(unknown_op.to_rust_op().is_err());
254    }
255    
256    #[test]
257    fn test_from_python_string() {
258        assert_eq!(BinOps::from_python_string("Add"), Some(BinOps::Add));
259        assert_eq!(BinOps::from_python_string("Unknown"), None);
260        assert_eq!(BinOps::parse_or_unknown("Invalid"), BinOps::Unknown);
261    }
262}
python_ast/ast/tree/bin_ops.rs

python_ast/ast/tree/
bin_ops.rs
