python_ast/ast/tree/
async_for.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, 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, 'py> FromPyObject<'a, 'py> for AsyncFor {
30    type Error = pyo3::PyErr;
31    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
32        // Extract target
33        let target: ExprType = ob.getattr("target")?.extract()?;
34        
35        // Extract iter
36        let iter: ExprType = ob.getattr("iter")?.extract()?;
37        
38        // Extract body
39        let body: Vec<Statement> = extract_list(&ob, "body", "async for body")?;
40        
41        // Extract orelse (optional)
42        let orelse: Vec<Statement> = extract_list(&ob, "orelse", "async for orelse").unwrap_or_default();
43        
44        Ok(AsyncFor {
45            target,
46            iter,
47            body,
48            orelse,
49            lineno: ob.lineno(),
50            col_offset: ob.col_offset(),
51            end_lineno: ob.end_lineno(),
52            end_col_offset: ob.end_col_offset(),
53        })
54    }
55}
56
57impl Node for AsyncFor {
58    fn lineno(&self) -> Option<usize> { self.lineno }
59    fn col_offset(&self) -> Option<usize> { self.col_offset }
60    fn end_lineno(&self) -> Option<usize> { self.end_lineno }
61    fn end_col_offset(&self) -> Option<usize> { self.end_col_offset }
62}
63
64impl CodeGen for AsyncFor {
65    type Context = CodeGenContext;
66    type Options = PythonOptions;
67    type SymbolTable = SymbolTableScopes;
68
69    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
70        // Process target, iter, body, and orelse
71        let symbols = self.target.find_symbols(symbols);
72        let symbols = self.iter.find_symbols(symbols);
73        let symbols = self.body.into_iter().fold(symbols, |acc, stmt| stmt.find_symbols(acc));
74        self.orelse.into_iter().fold(symbols, |acc, stmt| stmt.find_symbols(acc))
75    }
76
77    fn to_rust(
78        self,
79        ctx: Self::Context,
80        options: Self::Options,
81        symbols: Self::SymbolTable,
82    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
83        let target = self.target.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
84        let iter_expr = self.iter.to_rust(ctx.clone(), options.clone(), symbols.clone())?;
85
86        let has_else = !self.orelse.is_empty();
87        // Break-tracking is only needed when a break belonging to this loop
88        // could skip the else clause.
89        let tracks_break = has_else && crate::loop_body_has_direct_break(&self.body);
90        let body_ctx = crate::CodeGenContext::Loop {
91            has_else: tracks_break,
92            parent: Box::new(ctx.clone()),
93        };
94        let body_tokens: Result<Vec<TokenStream>, Box<dyn std::error::Error>> = self.body.into_iter()
95            .map(|stmt| stmt.to_rust(body_ctx.clone(), options.clone(), symbols.clone()))
96            .collect();
97        let body_tokens = body_tokens?;
98
99        // Full `async for` semantics need an async-stream protocol; until
100        // then, iterate the expression synchronously. This preserves the
101        // loop's iteration and target binding (previously the body ran once
102        // and the iterator was discarded entirely).
103        if !has_else {
104            Ok(quote! {
105                for #target in #iter_expr {
106                    #(#body_tokens;)*
107                }
108            })
109        } else {
110            let else_body_tokens: Result<Vec<TokenStream>, Box<dyn std::error::Error>> = self.orelse.into_iter()
111                .map(|stmt| stmt.to_rust(ctx.clone(), options.clone(), symbols.clone()))
112                .collect();
113            let else_body_tokens = else_body_tokens?;
114            if tracks_break {
115                Ok(quote! {
116                    {
117                        let mut __rython_broke = false;
118                        for #target in #iter_expr {
119                            #(#body_tokens;)*
120                        }
121                        if !__rython_broke {
122                            #(#else_body_tokens;)*
123                        }
124                    }
125                })
126            } else {
127                // No break can skip the else clause; run it unconditionally.
128                Ok(quote! {
129                    {
130                        for #target in #iter_expr {
131                            #(#body_tokens;)*
132                        }
133                        #(#else_body_tokens;)*
134                    }
135                })
136            }
137        }
138    }
139}
140
141#[cfg(test)]
142mod tests {
143    // Tests would go here - currently commented out as they need full AST infrastructure
144    // create_parse_test!(test_simple_async_for, "async for item in async_iter:\n    pass", "test.py");
145}
python_ast/ast/tree/async_for.rs

python_ast/ast/tree/
async_for.rs
