diff --git a/crates/py2erg/convert.rs b/crates/py2erg/convert.rs index 00fc451..0ed14fe 100644 --- a/crates/py2erg/convert.rs +++ b/crates/py2erg/convert.rs @@ -2450,6 +2450,33 @@ impl ASTConverter { Expr::ClassDef(classdef) } + fn convert_for(&mut self, for_: py_ast::StmtFor) -> Expr { + let loc = for_.location(); + let iter = self.convert_expr(*for_.iter); + let if_block_id = self.block_id_counter + 1; + let block = self.convert_for_body(Some(*for_.target), for_.body); + let for_ident = self.convert_ident("for".to_string(), loc); + let for_acc = Expr::Accessor(Accessor::Ident(for_ident)); + if for_.orelse.is_empty() { + for_acc.call2(iter, Expr::Lambda(block)) + } else { + let else_block = self.convert_block(for_.orelse, BlockKind::Else { if_block_id }); + let params = Params::empty(); + let sig = LambdaSignature::new(params, None, TypeBoundSpecs::empty()); + let op = Token::from_str(TokenKind::FuncArrow, "->"); + let else_block = Lambda::new(sig, op, else_block, DefId(0)); + let args = Args::pos_only( + vec![ + PosArg::new(iter), + PosArg::new(Expr::Lambda(block)), + PosArg::new(Expr::Lambda(else_block)), + ], + None, + ); + for_acc.call_expr(args) + } + } + fn get_t_spec(&self, name: &str) -> Option<&PyTypeSpec> { if self.contexts.len() == 1 { self.pyi_types.get_type(name) @@ -2817,32 +2844,55 @@ impl ASTConverter { self.convert_funcdef(func_def, true) } py_ast::Stmt::ClassDef(class_def) => self.convert_classdef(class_def), - py_ast::Stmt::For(for_) => { - let loc = for_.location(); - let iter = self.convert_expr(*for_.iter); - let block = self.convert_for_body(Some(*for_.target), for_.body); - let for_ident = self.convert_ident("for".to_string(), loc); - let for_acc = Expr::Accessor(Accessor::Ident(for_ident)); - for_acc.call2(iter, Expr::Lambda(block)) - } + py_ast::Stmt::For(for_) => self.convert_for(for_), py_ast::Stmt::AsyncFor(for_) => { - let loc = for_.location(); - let iter = self.convert_expr(*for_.iter); - let block = self.convert_for_body(Some(*for_.target), for_.body); - let for_ident = self.convert_ident("for".to_string(), loc); - let for_acc = Expr::Accessor(Accessor::Ident(for_ident)); - for_acc.call2(iter, Expr::Lambda(block)) + let py_ast::StmtAsyncFor { + target, + iter, + body, + orelse, + range, + type_comment, + } = for_; + let for_ = py_ast::StmtFor { + target, + iter, + body, + orelse, + range, + type_comment, + }; + self.convert_for(for_) } py_ast::Stmt::While(while_) => { let loc = while_.location(); let test = self.convert_expr(*while_.test); let params = Params::empty(); let empty_sig = LambdaSignature::new(params, None, TypeBoundSpecs::empty()); + let if_block_id = self.block_id_counter + 1; let block = self.convert_block(while_.body, BlockKind::While); let body = Lambda::new(empty_sig, Token::DUMMY, block, DefId(0)); let while_ident = self.convert_ident("while".to_string(), loc); let while_acc = Expr::Accessor(Accessor::Ident(while_ident)); - while_acc.call2(test, Expr::Lambda(body)) + if while_.orelse.is_empty() { + while_acc.call2(test, Expr::Lambda(body)) + } else { + let else_block = + self.convert_block(while_.orelse, BlockKind::Else { if_block_id }); + let params = Params::empty(); + let sig = LambdaSignature::new(params, None, TypeBoundSpecs::empty()); + let op = Token::from_str(TokenKind::FuncArrow, "->"); + let else_body = Lambda::new(sig, op, else_block, DefId(0)); + let args = Args::pos_only( + vec![ + PosArg::new(test), + PosArg::new(Expr::Lambda(body)), + PosArg::new(Expr::Lambda(else_body)), + ], + None, + ); + while_acc.call_expr(args) + } } py_ast::Stmt::If(if_) => { let loc = if_.location(); @@ -2958,8 +3008,13 @@ impl ASTConverter { } py_ast::Stmt::Try(try_) => { let if_block_id = self.block_id_counter + 1; - let chunks = self.convert_block(try_.body, BlockKind::Try).into_iter(); + let mut chunks = self.convert_block(try_.body, BlockKind::Try); + for py_ast::ExceptHandler::ExceptHandler(handler) in try_.handlers { + chunks + .extend(self.convert_block(handler.body, BlockKind::Else { if_block_id })); + } let dummy = chunks + .into_iter() .chain(self.convert_block(try_.orelse, BlockKind::Else { if_block_id })) .chain(self.convert_block(try_.finalbody, BlockKind::Else { if_block_id })) .collect();