python_ast/ast/tree/
aug_assign.rs

1use proc_macro2::TokenStream;
2use pyo3::{Bound, 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> FromPyObject<'a> for AugAssign {
28    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
29        // Extract target
30        let target = ob.extract_attr_with_context("target", "augmented assignment target")?;
31        let target: ExprType = target.extract()?;
32        
33        // Extract operator
34        let op = ob.extract_attr_with_context("op", "augmented assignment operator")?;
35        let op_type_str = op.extract_type_name("augmented assignment operator")?;
36        let op = BinOps::parse_or_unknown(&op_type_str);
37        
38        // Extract value
39        let value = ob.extract_attr_with_context("value", "augmented assignment value")?;
40        let value: ExprType = value.extract()?;
41        
42        Ok(AugAssign {
43            target,
44            op,
45            value,
46            lineno: ob.lineno(),
47            col_offset: ob.col_offset(),
48            end_lineno: ob.end_lineno(),
49            end_col_offset: ob.end_col_offset(),
50        })
51    }
52}
53
54impl Node for AugAssign {
55    fn lineno(&self) -> Option<usize> { self.lineno }
56    fn col_offset(&self) -> Option<usize> { self.col_offset }
57    fn end_lineno(&self) -> Option<usize> { self.end_lineno }
58    fn end_col_offset(&self) -> Option<usize> { self.end_col_offset }
59}
60
61impl CodeGen for AugAssign {
62    type Context = CodeGenContext;
63    type Options = PythonOptions;
64    type SymbolTable = SymbolTableScopes;
65
66    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
67        // Process the value for symbols, but don't add new symbols for augmented assignment
68        self.value.find_symbols(symbols)
69    }
70
71    fn to_rust(
72        self,
73        ctx: Self::Context,
74        options: Self::Options,
75        symbols: Self::SymbolTable,
76    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
77        let target = self.target.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
78        let value = self.value.to_rust(ctx, options, symbols)?;
79        
80        // Generate the appropriate augmented assignment operator
81        match self.op {
82            BinOps::Add => Ok(quote!(#target += #value)),
83            BinOps::Sub => Ok(quote!(#target -= #value)),
84            BinOps::Mult => Ok(quote!(#target *= #value)),
85            BinOps::Div => Ok(quote!(#target /= #value)),
86            BinOps::FloorDiv => Ok(quote!(#target /= #value)), // Rust doesn't have floor div assign
87            BinOps::Mod => Ok(quote!(#target %= #value)),
88            BinOps::BitAnd => Ok(quote!(#target &= #value)),
89            BinOps::BitOr => Ok(quote!(#target |= #value)),
90            BinOps::BitXor => Ok(quote!(#target ^= #value)),
91            BinOps::LShift => Ok(quote!(#target <<= #value)),
92            BinOps::RShift => Ok(quote!(#target >>= #value)),
93            BinOps::Pow => {
94                // Rust doesn't have **= operator, so we need to expand it
95                Ok(quote!(#target = (#target).pow(#value)))
96            },
97            BinOps::MatMult => {
98                // Matrix multiplication assignment - not directly supported in Rust
99                // Would need specific matrix library support
100                Err(format!("Matrix multiplication assignment not supported in Rust").into())
101            },
102            BinOps::Unknown => {
103                Err(format!("Unknown augmented assignment operator").into())
104            },
105        }
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112    use crate::create_parse_test;
113
114    create_parse_test!(test_add_assign, "x += 1", "test.py");
115    create_parse_test!(test_sub_assign, "x -= 1", "test.py");
116    create_parse_test!(test_mul_assign, "x *= 2", "test.py");
117    create_parse_test!(test_div_assign, "x /= 3", "test.py");
118    create_parse_test!(test_mod_assign, "x %= 4", "test.py");
119    create_parse_test!(test_pow_assign, "x **= 2", "test.py");
120    create_parse_test!(test_bitand_assign, "x &= 5", "test.py");
121    create_parse_test!(test_bitor_assign, "x |= 6", "test.py");
122    create_parse_test!(test_bitxor_assign, "x ^= 7", "test.py");
123    create_parse_test!(test_lshift_assign, "x <<= 2", "test.py");
124    create_parse_test!(test_rshift_assign, "x >>= 3", "test.py");
125}
python_ast/ast/tree/aug_assign.rs

python_ast/ast/tree/
aug_assign.rs
