python_ast/ast/tree/
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, SymbolTableNode,
8    SymbolTableScopes,
9};
10
11#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
12pub struct Assign {
13    pub targets: Vec<ExprType>,
14    pub value: ExprType,
15    pub type_comment: Option<String>,
16}
17
18impl<'a> FromPyObject<'a> for Assign {
19    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
20        let targets: Vec<ExprType> = ob
21            .getattr("targets")
22            .expect(
23                ob.error_message("<unknown>", "error getting assignment targets")
24                    .as_str(),
25            )
26            .extract()
27            .expect("extracting assignment targets");
28
29        let python_value = ob.getattr("value").expect(
30            ob.error_message("<unknown>", "assignment statement value not found")
31                .as_str(),
32        );
33
34        let value = python_value.extract().expect(
35            ob.error_message("<unknown>", "error getting value of assignment statement")
36                .as_str(),
37        );
38
39        Ok(Assign {
40            targets: targets,
41            value: value,
42            type_comment: None,
43        })
44    }
45}
46
47impl<'a> CodeGen for Assign {
48    type Context = CodeGenContext;
49    type Options = PythonOptions;
50    type SymbolTable = SymbolTableScopes;
51
52    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
53        let mut symbols = symbols;
54        let mut position = 0;
55        for target in self.targets {
56            // Only add symbols for Name assignments, not for Attribute assignments
57            if let ExprType::Name(name) = target {
58                symbols.insert(
59                    name.id,
60                    SymbolTableNode::Assign {
61                        position: position,
62                        value: self.value.clone(),
63                    },
64                );
65            }
66            // Could also handle other target types here if needed
67            position += 1;
68        }
69        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        let mut target_streams = Vec::new();
79        
80        // Convert each target to Rust code
81        for target in self.targets {
82            let target_code = target.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
83            target_streams.push(target_code);
84        }
85        
86        let value = self.value.to_rust(ctx, options, symbols)?;
87        
88        // For single target assignment
89        if target_streams.len() == 1 {
90            let target = &target_streams[0];
91            // Check if this is a new variable declaration or reassignment
92            // For now, we'll use `let` for new declarations
93            Ok(quote!(let #target = #value;))
94        } else {
95            // For multiple assignment targets like: a, b = 1, 2
96            // Use tuple destructuring in Rust
97            Ok(quote! {
98                let (#(#target_streams),*) = #value;
99            })
100        }
101    }
102}
python_ast/ast/tree/assign.rs

python_ast/ast/tree/
assign.rs
