python_ast/ast/tree/
async_for.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, Statement, SymbolTableScopes,
8    extract_list,
9};
10
11/// Async for loop (async for target in iter: ...)
12#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
13pub struct AsyncFor {
14    /// The target variable(s)
15    pub target: ExprType,
16    /// The iterable expression
17    pub iter: ExprType,
18    /// The body of the loop
19    pub body: Vec<Statement>,
20    /// The else clause (executed when the loop completes normally)
21    pub orelse: Vec<Statement>,
22    /// Position information
23    pub lineno: Option<usize>,
24    pub col_offset: Option<usize>,
25    pub end_lineno: Option<usize>,
26    pub end_col_offset: Option<usize>,
27}
28
29impl<'a> FromPyObject<'a> for AsyncFor {
30    fn extract_bound(ob: &Bound<'a, PyAny>) -> PyResult<Self> {
31        // Extract target
32        let target: ExprType = ob.getattr("target")?.extract()?;
33        
34        // Extract iter
35        let iter: ExprType = ob.getattr("iter")?.extract()?;
36        
37        // Extract body
38        let body: Vec<Statement> = extract_list(ob, "body", "async for body")?;
39        
40        // Extract orelse (optional)
41        let orelse: Vec<Statement> = extract_list(ob, "orelse", "async for orelse").unwrap_or_default();
42        
43        Ok(AsyncFor {
44            target,
45            iter,
46            body,
47            orelse,
48            lineno: ob.lineno(),
49            col_offset: ob.col_offset(),
50            end_lineno: ob.end_lineno(),
51            end_col_offset: ob.end_col_offset(),
52        })
53    }
54}
55
56impl Node for AsyncFor {
57    fn lineno(&self) -> Option<usize> { self.lineno }
58    fn col_offset(&self) -> Option<usize> { self.col_offset }
59    fn end_lineno(&self) -> Option<usize> { self.end_lineno }
60    fn end_col_offset(&self) -> Option<usize> { self.end_col_offset }
61}
62
63impl CodeGen for AsyncFor {
64    type Context = CodeGenContext;
65    type Options = PythonOptions;
66    type SymbolTable = SymbolTableScopes;
67
68    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
69        // Process target, iter, body, and orelse
70        let symbols = self.target.find_symbols(symbols);
71        let symbols = self.iter.find_symbols(symbols);
72        let symbols = self.body.into_iter().fold(symbols, |acc, stmt| stmt.find_symbols(acc));
73        self.orelse.into_iter().fold(symbols, |acc, stmt| stmt.find_symbols(acc))
74    }
75
76    fn to_rust(
77        self,
78        ctx: Self::Context,
79        options: Self::Options,
80        symbols: Self::SymbolTable,
81    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
82        // Generate iter expression
83        let _iter_expr = self.iter.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
84        
85        // Generate body
86        let body_tokens: Result<Vec<TokenStream>, Box<dyn std::error::Error>> = self.body.into_iter()
87            .map(|stmt| stmt.to_rust(ctx.clone(), options.clone(), symbols.clone()))
88            .collect();
89        let body_tokens = body_tokens?;
90
91        // Generate else clause if present
92        let else_tokens = if !self.orelse.is_empty() {
93            let else_body_tokens: Result<Vec<TokenStream>, Box<dyn std::error::Error>> = self.orelse.into_iter()
94                .map(|stmt| stmt.to_rust(ctx.clone(), options.clone(), symbols.clone()))
95                .collect();
96            let else_body_tokens = else_body_tokens?;
97            quote! {
98                // Else clause (executed when loop completes normally)
99                #(#else_body_tokens)*
100            }
101        } else {
102            quote!()
103        };
104
105        // For now, generate a simplified async iteration
106        // In practice, this would need proper async stream handling
107        Ok(quote! {
108            {
109                // Async for loop - simplified translation  
110                // Python's async for doesn't map directly to Rust's async streams
111                // This would typically use futures::stream::StreamExt
112                #(#body_tokens)*
113                
114                // Else clause
115                #else_tokens
116            }
117        })
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    // Tests would go here - currently commented out as they need full AST infrastructure
124    // create_parse_test!(test_simple_async_for, "async for item in async_iter:\n    pass", "test.py");
125}