use rustpython_parser::ast::{
BooleanOperator, Comparison, Comprehension, ComprehensionKind, ExpressionType, Keyword,
Located, Number, Operator, Parameter, Parameters, StringGroup, UnaryOperator, Varargs,
};
fn clone_number(num: &Number) -> Number {
match num {
Number::Integer { value } => Number::Integer {
value: value.clone(),
},
Number::Float { value } => Number::Float { value: *value },
Number::Complex { real, imag } => Number::Complex {
real: *real,
imag: *imag,
},
}
}
fn clone_string_group(group: &StringGroup) -> StringGroup {
match group {
StringGroup::Constant { value } => StringGroup::Constant {
value: value.clone(),
},
StringGroup::FormattedValue {
value,
conversion,
spec,
} => StringGroup::FormattedValue {
value: Box::new(clone_loc_expr(value)),
conversion: *conversion,
spec: spec.as_deref().map(|sp| Box::new(clone_string_group(sp))),
},
StringGroup::Joined { values } => StringGroup::Joined {
values: values.iter().map(clone_string_group).collect::<Vec<_>>(),
},
}
}
fn clone_unary_op(op: &UnaryOperator) -> UnaryOperator {
match op {
UnaryOperator::Not => UnaryOperator::Not,
UnaryOperator::Inv => UnaryOperator::Inv,
UnaryOperator::Pos => UnaryOperator::Pos,
UnaryOperator::Neg => UnaryOperator::Neg,
}
}
fn clone_bin_op(op: &Operator) -> Operator {
match op {
Operator::Add => Operator::Add,
Operator::Sub => Operator::Sub,
Operator::Mult => Operator::Mult,
Operator::MatMult => Operator::MatMult,
Operator::Div => Operator::Div,
Operator::Mod => Operator::Mod,
Operator::Pow => Operator::Pow,
Operator::LShift => Operator::LShift,
Operator::RShift => Operator::RShift,
Operator::BitOr => Operator::BitOr,
Operator::BitXor => Operator::BitXor,
Operator::BitAnd => Operator::BitAnd,
Operator::FloorDiv => Operator::FloorDiv,
}
}
fn clone_comp_op(op: &Comparison) -> Comparison {
match op {
Comparison::Equal => Comparison::Equal,
Comparison::NotEqual => Comparison::NotEqual,
Comparison::Less => Comparison::Less,
Comparison::LessOrEqual => Comparison::LessOrEqual,
Comparison::Greater => Comparison::Greater,
Comparison::GreaterOrEqual => Comparison::GreaterOrEqual,
Comparison::Is => Comparison::Is,
Comparison::IsNot => Comparison::IsNot,
Comparison::In => Comparison::In,
Comparison::NotIn => Comparison::NotIn,
}
}
fn clone_bool_op(op: &BooleanOperator) -> BooleanOperator {
match op {
BooleanOperator::And => BooleanOperator::And,
BooleanOperator::Or => BooleanOperator::Or,
}
}
fn clone_param(param: &Parameter) -> Parameter {
Parameter {
location: param.location,
arg: param.arg.clone(),
annotation: param
.annotation
.as_deref()
.map(|a| Box::new(clone_loc_expr(a))),
}
}
fn clone_varargs(varargs: &Varargs) -> Varargs {
match varargs {
Varargs::None => Varargs::None,
Varargs::Unnamed => Varargs::Unnamed,
Varargs::Named(name) => Varargs::Named(clone_param(name)),
}
}
fn clone_params(params: &Parameters) -> Parameters {
Parameters {
posonlyargs_count: params.posonlyargs_count,
args: params.args.iter().map(clone_param).collect::<Vec<_>>(),
vararg: clone_varargs(¶ms.vararg),
kwonlyargs: params
.kwonlyargs
.iter()
.map(clone_param)
.collect::<Vec<_>>(),
kw_defaults: params
.kw_defaults
.iter()
.map(|def| def.as_ref().map(clone_loc_expr))
.collect::<Vec<_>>(),
kwarg: clone_varargs(¶ms.kwarg),
defaults: params
.defaults
.iter()
.map(clone_loc_expr)
.collect::<Vec<_>>(),
}
}
fn clone_kw(keyword: &Keyword) -> Keyword {
Keyword {
name: keyword.name.clone(),
value: clone_loc_expr(&keyword.value),
}
}
fn clone_comprehension_kind(kind: &ComprehensionKind) -> ComprehensionKind {
match kind {
ComprehensionKind::Dict { key, value } => ComprehensionKind::Dict {
key: clone_loc_expr(key),
value: clone_loc_expr(value),
},
ComprehensionKind::List { element } => ComprehensionKind::List {
element: clone_loc_expr(element),
},
ComprehensionKind::Set { element } => ComprehensionKind::Set {
element: clone_loc_expr(element),
},
ComprehensionKind::GeneratorExpression { element } => {
ComprehensionKind::GeneratorExpression {
element: clone_loc_expr(element),
}
}
}
}
pub fn clone_loc_expr(expr: &Located<ExpressionType>) -> Located<ExpressionType> {
Located {
node: clone_expr(&expr.node),
location: expr.location,
}
}
pub fn clone_expr(expr: &ExpressionType) -> ExpressionType {
match expr {
ExpressionType::None => ExpressionType::None,
ExpressionType::Ellipsis => ExpressionType::Ellipsis,
ExpressionType::True => ExpressionType::True,
ExpressionType::False => ExpressionType::False,
ExpressionType::Identifier { name } => ExpressionType::Identifier { name: name.clone() },
ExpressionType::Number { value } => ExpressionType::Number {
value: clone_number(value),
},
ExpressionType::String { value } => ExpressionType::String {
value: clone_string_group(value),
},
ExpressionType::Attribute { value, name } => ExpressionType::Attribute {
value: Box::new(clone_loc_expr(value)),
name: name.clone(),
},
ExpressionType::Subscript { a, b } => ExpressionType::Subscript {
a: Box::new(clone_loc_expr(a)),
b: Box::new(clone_loc_expr(b)),
},
ExpressionType::Slice { elements } => ExpressionType::Slice {
elements: elements.iter().map(clone_loc_expr).collect::<Vec<_>>(),
},
ExpressionType::Bytes { value } => ExpressionType::Bytes {
value: value.clone(),
},
ExpressionType::Call {
function,
args,
keywords,
} => ExpressionType::Call {
function: Box::new(clone_loc_expr(function)),
args: args.iter().map(clone_loc_expr).collect::<Vec<_>>(),
keywords: keywords.iter().map(clone_kw).collect::<Vec<_>>(),
},
ExpressionType::Unop { op, a } => ExpressionType::Unop {
op: clone_unary_op(op),
a: Box::new(clone_loc_expr(a)),
},
ExpressionType::Binop { a, op, b } => ExpressionType::Binop {
a: Box::new(clone_loc_expr(a)),
op: clone_bin_op(op),
b: Box::new(clone_loc_expr(b)),
},
ExpressionType::Compare { vals, ops } => ExpressionType::Compare {
ops: ops.iter().map(clone_comp_op).collect::<Vec<_>>(),
vals: vals.iter().map(clone_loc_expr).collect::<Vec<_>>(),
},
ExpressionType::BoolOp { op, values } => ExpressionType::BoolOp {
op: clone_bool_op(op),
values: values.iter().map(clone_loc_expr).collect::<Vec<_>>(),
},
ExpressionType::Lambda { args, body } => ExpressionType::Lambda {
args: Box::new(clone_params(args)),
body: Box::new(clone_loc_expr(body)),
},
ExpressionType::IfExpression { test, body, orelse } => ExpressionType::IfExpression {
test: Box::new(clone_loc_expr(test)),
body: Box::new(clone_loc_expr(body)),
orelse: Box::new(clone_loc_expr(orelse)),
},
ExpressionType::Dict { elements } => ExpressionType::Dict {
elements: elements
.iter()
.map(|(key, value)| (key.as_ref().map(clone_loc_expr), clone_loc_expr(value)))
.collect::<Vec<_>>(),
},
ExpressionType::Set { elements } => ExpressionType::Set {
elements: elements.iter().map(clone_loc_expr).collect::<Vec<_>>(),
},
ExpressionType::List { elements } => ExpressionType::List {
elements: elements.iter().map(clone_loc_expr).collect::<Vec<_>>(),
},
ExpressionType::Tuple { elements } => ExpressionType::Tuple {
elements: elements.iter().map(clone_loc_expr).collect::<Vec<_>>(),
},
ExpressionType::Yield { value } => ExpressionType::Yield {
value: value.as_ref().map(|val| Box::new(clone_loc_expr(val))),
},
ExpressionType::YieldFrom { value } => ExpressionType::YieldFrom {
value: Box::new(clone_loc_expr(value)),
},
ExpressionType::Await { value } => ExpressionType::Await {
value: Box::new(clone_loc_expr(value)),
},
ExpressionType::NamedExpression { left, right } => ExpressionType::NamedExpression {
left: Box::new(clone_loc_expr(left)),
right: Box::new(clone_loc_expr(right)),
},
ExpressionType::Starred { value } => ExpressionType::Starred {
value: Box::new(clone_loc_expr(value)),
},
ExpressionType::Comprehension { kind, generators } => ExpressionType::Comprehension {
kind: Box::new(clone_comprehension_kind(kind)),
generators: generators
.iter()
.map(|gen| Comprehension {
location: gen.location,
target: clone_loc_expr(&gen.target),
iter: clone_loc_expr(&gen.iter),
ifs: gen.ifs.iter().map(clone_loc_expr).collect::<Vec<_>>(),
is_async: gen.is_async,
})
.collect::<Vec<_>>(),
},
}
}