use super::*;
use rustpython_compiler_core::SourceFile;
pub(super) fn ast_to_object(
clause: ast::ElifElseClause,
mut rest: alloc::vec::IntoIter<ast::ElifElseClause>,
vm: &VirtualMachine,
source_file: &SourceFile,
) -> PyObjectRef {
let ast::ElifElseClause {
node_index: _,
range,
test,
body,
} = clause;
let Some(test) = test else {
assert!(rest.len() == 0);
return body.ast_to_object(vm, source_file);
};
let node = NodeAst
.into_ref_with_type(vm, pyast::NodeStmtIf::static_type().to_owned())
.unwrap();
let dict = node.as_object().dict().unwrap();
dict.set_item("test", test.ast_to_object(vm, source_file), vm)
.unwrap();
dict.set_item("body", body.ast_to_object(vm, source_file), vm)
.unwrap();
let orelse = if let Some(next) = rest.next() {
if next.test.is_some() {
let next = ast::ElifElseClause {
range: TextRange::new(next.range.start(), range.end()),
..next
};
vm.ctx
.new_list(vec![ast_to_object(next, rest, vm, source_file)])
.into()
} else {
next.body.ast_to_object(vm, source_file)
}
} else {
vm.ctx.new_list(vec![]).into()
};
dict.set_item("orelse", orelse, vm).unwrap();
node_add_location(&dict, range, vm, source_file);
node.into()
}
pub(super) fn ast_from_object(
vm: &VirtualMachine,
source_file: &SourceFile,
object: PyObjectRef,
) -> PyResult<ast::StmtIf> {
let test = Node::ast_from_object(vm, source_file, get_node_field(vm, &object, "test", "If")?)?;
let body = Node::ast_from_object(vm, source_file, get_node_field(vm, &object, "body", "If")?)?;
let orelse: Vec<ast::Stmt> = Node::ast_from_object(
vm,
source_file,
get_node_field(vm, &object, "orelse", "If")?,
)?;
let range = range_from_object(vm, source_file, object, "If")?;
let elif_else_clauses = if orelse.is_empty() {
vec![]
} else if let [ast::Stmt::If(_)] = &*orelse {
let Some(ast::Stmt::If(ast::StmtIf {
node_index: _,
range,
test,
body,
mut elif_else_clauses,
})) = orelse.into_iter().next()
else {
unreachable!()
};
elif_else_clauses.insert(
0,
ast::ElifElseClause {
node_index: Default::default(),
range,
test: Some(*test),
body,
},
);
elif_else_clauses
} else {
vec![ast::ElifElseClause {
node_index: Default::default(),
range,
test: None,
body: orelse,
}]
};
Ok(ast::StmtIf {
node_index: Default::default(),
test,
body,
elif_else_clauses,
range,
})
}