python_ast/ast/tree/
aug_assign.rs

1use proc_macro2::TokenStream;
2use pyo3::{Borrowed, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods};
3use quote::quote;
4use serde::{Deserialize, Serialize};
5
6use crate::{
7    CodeGen, CodeGenContext, ExprType, Node, PythonOptions, SymbolTableScopes,
8    BinOps, FromPythonString, PyAttributeExtractor,
9};
10
11/// Augmented assignment statement (e.g., x += 1, y -= 2, etc.)
12#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
13pub struct AugAssign {
14    /// The target being assigned to (left side)
15    pub target: ExprType,
16    /// The operator (+=, -=, *=, etc.)
17    pub op: BinOps,
18    /// The value being assigned (right side)
19    pub value: ExprType,
20    /// Position information
21    pub lineno: Option<usize>,
22    pub col_offset: Option<usize>,
23    pub end_lineno: Option<usize>,
24    pub end_col_offset: Option<usize>,
25}
26
27impl<'a, 'py> FromPyObject<'a, 'py> for AugAssign {
28    type Error = pyo3::PyErr;
29    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
30        // Extract target
31        let target = ob.extract_attr_with_context("target", "augmented assignment target")?;
32        let target: ExprType = target.extract()?;
33        
34        // Extract operator
35        let op = ob.extract_attr_with_context("op", "augmented assignment operator")?;
36        let op_type_str = op.extract_type_name("augmented assignment operator")?;
37        let op = BinOps::parse_or_unknown(&op_type_str);
38        
39        // Extract value
40        let value = ob.extract_attr_with_context("value", "augmented assignment value")?;
41        let value: ExprType = value.extract()?;
42        
43        Ok(AugAssign {
44            target,
45            op,
46            value,
47            lineno: ob.lineno(),
48            col_offset: ob.col_offset(),
49            end_lineno: ob.end_lineno(),
50            end_col_offset: ob.end_col_offset(),
51        })
52    }
53}
54
55impl Node for AugAssign {
56    fn lineno(&self) -> Option<usize> { self.lineno }
57    fn col_offset(&self) -> Option<usize> { self.col_offset }
58    fn end_lineno(&self) -> Option<usize> { self.end_lineno }
59    fn end_col_offset(&self) -> Option<usize> { self.end_col_offset }
60}
61
62impl CodeGen for AugAssign {
63    type Context = CodeGenContext;
64    type Options = PythonOptions;
65    type SymbolTable = SymbolTableScopes;
66
67    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
68        // Process the value for symbols, but don't add new symbols for augmented assignment
69        self.value.find_symbols(symbols)
70    }
71
72    fn to_rust(
73        self,
74        ctx: Self::Context,
75        options: Self::Options,
76        symbols: Self::SymbolTable,
77    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
78        // Compound assignment to a subscript (`counts[k] += 1`) is a
79        // read-modify-write: the Load lowering of the target is a cloned
80        // temporary (py_index), not a place, so read via py_index, combine,
81        // and store back via py_set_index. The index is evaluated once.
82        if let ExprType::Subscript(sub) = &self.target {
83            // The receiver must be a place (see subscript_receiver_place);
84            // a cloned receiver would silently drop the write-back.
85            let receiver = crate::subscript_receiver_place(
86                &sub.value,
87                ctx.clone(),
88                options.clone(),
89                symbols.clone(),
90            )?;
91            let index = match &sub.kind {
92                crate::SubscriptKind::Index(index) => index
93                    .clone()
94                    .to_rust(ctx.clone(), options.clone(), symbols.clone())?,
95                crate::SubscriptKind::Slice { .. } => {
96                    return Err(
97                        "augmented assignment to a slice (`x[a:b] += ...`) is not supported"
98                            .to_string()
99                            .into(),
100                    )
101                }
102            };
103            let value = self.value.to_rust(ctx, options, symbols)?;
104            let elem = quote!(__rython_elem);
105            let combined = combine_op(&self.op, &elem, &value)?;
106            // The receiver place is bound once so a nested chain
107            // (`grid[i][j] += 1`) evaluates its intermediate lookups — and
108            // any side effects in their indices — exactly once.
109            return Ok(quote! {
110                {
111                    let __rython_recv = &mut (#receiver);
112                    let __rython_idx = #index;
113                    let __rython_elem = (__rython_recv).py_index(__rython_idx.clone())?;
114                    (__rython_recv).py_set_index(__rython_idx, #combined)?;
115                }
116            });
117        }
118
119        let target = self.target.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
120        let value = self.value.to_rust(ctx, options, symbols)?;
121
122        // Generate the appropriate augmented assignment operator
123        match self.op {
124            // `+=` mirrors Python's `+` (string concat, list concat,
125            // numeric promotion) via PyAdd.
126            BinOps::Add => Ok(quote!(#target = (#target).py_add(&(#value)))),
127            BinOps::Sub => Ok(quote!(#target -= #value)),
128            BinOps::Mult => Ok(quote!(#target *= #value)),
129            // Python's `/` is TRUE division: `x /= 2` on an int yields a
130            // float. Rust's `/=` on an integer truncates silently, so
131            // mirror the BinOp lowering instead (an int target then fails
132            // to compile, which is loud rather than quietly wrong).
133            BinOps::Div => Ok(quote!(#target = (#target) as f64 / (#value) as f64)),
134            // Python // and % floor toward negative infinity / take the
135            // divisor's sign; use the stdpython helpers instead of Rust's
136            // truncating operators.
137            BinOps::FloorDiv => Ok(quote!(#target = py_floordiv(#target, #value))),
138            BinOps::Mod => Ok(quote!(#target = py_mod(#target, #value))),
139            BinOps::BitAnd => Ok(quote!(#target &= #value)),
140            BinOps::BitOr => Ok(quote!(#target |= #value)),
141            BinOps::BitXor => Ok(quote!(#target ^= #value)),
142            BinOps::LShift => Ok(quote!(#target <<= #value)),
143            BinOps::RShift => Ok(quote!(#target >>= #value)),
144            BinOps::Pow => {
145                // Rust doesn't have **= operator, so we need to expand it
146                Ok(quote!(#target = py_pow(#target, #value)))
147            },
148            BinOps::MatMult => {
149                // Matrix multiplication assignment - not directly supported in Rust
150                // Would need specific matrix library support
151                Err(format!("Matrix multiplication assignment not supported in Rust").into())
152            },
153            BinOps::Unknown => {
154                Err(format!("Unknown augmented assignment operator").into())
155            },
156        }
157    }
158}
159
160/// The read-modify-write combination for a compound assignment: how the
161/// current element and the operand produce the stored value.
162fn combine_op(
163    op: &BinOps,
164    elem: &TokenStream,
165    value: &TokenStream,
166) -> Result<TokenStream, Box<dyn std::error::Error>> {
167    Ok(match op {
168        BinOps::Add => quote!((#elem).py_add(&(#value))),
169        BinOps::Sub => quote!(#elem - #value),
170        BinOps::Mult => quote!(#elem * #value),
171        BinOps::Div => quote!((#elem) as f64 / (#value) as f64),
172        BinOps::FloorDiv => quote!(py_floordiv(#elem, #value)),
173        BinOps::Mod => quote!(py_mod(#elem, #value)),
174        BinOps::Pow => quote!(py_pow(#elem, #value)),
175        BinOps::BitAnd => quote!(#elem & #value),
176        BinOps::BitOr => quote!(#elem | #value),
177        BinOps::BitXor => quote!(#elem ^ #value),
178        BinOps::LShift => quote!(#elem << #value),
179        BinOps::RShift => quote!(#elem >> #value),
180        other => {
181            return Err(format!(
182                "augmented assignment operator {:?} not supported on subscripts",
183                other
184            )
185            .into())
186        }
187    })
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193    use crate::create_parse_test;
194
195    create_parse_test!(test_add_assign, "x += 1", "test.py");
196    create_parse_test!(test_sub_assign, "x -= 1", "test.py");
197    create_parse_test!(test_mul_assign, "x *= 2", "test.py");
198    create_parse_test!(test_div_assign, "x /= 3", "test.py");
199    create_parse_test!(test_mod_assign, "x %= 4", "test.py");
200    create_parse_test!(test_pow_assign, "x **= 2", "test.py");
201    create_parse_test!(test_bitand_assign, "x &= 5", "test.py");
202    create_parse_test!(test_bitor_assign, "x |= 6", "test.py");
203    create_parse_test!(test_bitxor_assign, "x ^= 7", "test.py");
204    create_parse_test!(test_lshift_assign, "x <<= 2", "test.py");
205    create_parse_test!(test_rshift_assign, "x >>= 3", "test.py");
206}
python_ast/ast/tree/aug_assign.rs

python_ast/ast/tree/
aug_assign.rs
