pylyzer/crates/py2erg/ast_util.rs
2023-05-10 00:06:10 +09:00

132 lines
4.8 KiB
Rust

use erg_common::log;
use rustpython_parser::ast::{
BooleanOperator, Comparison, ExpressionType, Keyword, Number, StringGroup,
};
pub fn number_to_string(num: &Number) -> String {
match num {
Number::Integer { value } => value.to_string(),
Number::Float { value } => value.to_string(),
Number::Complex { real, imag } => format!("{real}+{imag}j"),
}
}
pub fn keyword_length(keyword: &Keyword) -> usize {
if let Some(name) = &keyword.name {
name.len() + 1 + length(&keyword.value.node)
} else {
length(&keyword.value.node)
}
}
pub fn string_length(string: &StringGroup) -> usize {
match string {
StringGroup::Constant { value } => value.len(),
StringGroup::Joined { values } => values.iter().map(string_length).sum(),
other => {
log!(err "{other:?}");
0
}
}
}
pub fn comp_to_string(comp: &Comparison) -> String {
match comp {
Comparison::In => "in".to_string(),
Comparison::NotIn => "not in".to_string(),
Comparison::Is => "is".to_string(),
Comparison::IsNot => "is not".to_string(),
Comparison::Less => "<".to_string(),
Comparison::Greater => ">".to_string(),
Comparison::Equal => "==".to_string(),
Comparison::NotEqual => "!=".to_string(),
Comparison::LessOrEqual => "<=".to_string(),
Comparison::GreaterOrEqual => ">=".to_string(),
}
}
pub fn length(expr: &ExpressionType) -> usize {
match expr {
ExpressionType::Identifier { name } => name.len(),
ExpressionType::Number { value } => number_to_string(value).len(),
ExpressionType::String { value } => string_length(value),
ExpressionType::Attribute { value, name } => length(&value.node) + name.len() + 1,
ExpressionType::Subscript { a, b } => length(&a.node) + length(&b.node) + 2,
ExpressionType::Tuple { elements }
| ExpressionType::List { elements }
| ExpressionType::Set { elements } => {
if let (Some(first), Some(last)) = (elements.first(), elements.last()) {
2 + last.location.column() - first.location.column()
} else {
2
}
}
ExpressionType::Call {
function,
args,
keywords,
} => {
let args_len = if let (Some(first), Some(last)) = (args.first(), args.last()) {
last.location.column() - first.location.column()
} else {
0
};
let kw_len = if let (Some(first), Some(last)) = (keywords.first(), keywords.last()) {
last.value.location.column() - first.value.location.column()
} else {
0
};
length(&function.node) + args_len + kw_len + 2 // ()
}
ExpressionType::Unop { op: _, a } => 1 + length(&a.node),
ExpressionType::Binop { a, op: _, b } => length(&a.node) + 3 + length(&b.node),
ExpressionType::BoolOp { op, values } => match op {
BooleanOperator::And => values
.iter()
.map(|elem| length(&elem.node))
.fold(0, |acc, x| acc + x + 3),
BooleanOperator::Or => values
.iter()
.map(|elem| length(&elem.node))
.fold(0, |acc, x| acc + x + 2),
},
ExpressionType::Compare { vals, ops } => vals
.iter()
.zip(ops.iter())
.map(|(elem, op)| length(&elem.node) + comp_to_string(op).len())
.fold(0, |acc, x| acc + x + 2),
ExpressionType::IfExpression { test, body, orelse } => {
// x if y else z
length(&test.node) + 4 + length(&body.node) + 6 + length(&orelse.node)
}
ExpressionType::Lambda { args: _, body } => {
// lambda x: y
// TODO:
7 + 1 + length(&body.node)
}
ExpressionType::Await { value } => 5 + length(&value.node),
ExpressionType::Yield { value } => 5 + value.as_ref().map(|x| length(&x.node)).unwrap_or(0),
ExpressionType::NamedExpression { left, right } => {
// x := y
length(&left.node) + 4 + length(&right.node)
}
ExpressionType::Starred { value } => 1 + length(&value.node),
ExpressionType::False => 5,
ExpressionType::True | ExpressionType::None => 4,
ExpressionType::Ellipsis => 8,
other => {
log!(err "{other:?}");
0
}
}
}
pub fn accessor_name(expr: ExpressionType) -> Option<String> {
match expr {
ExpressionType::Identifier { name } => Some(name),
ExpressionType::Attribute { value, name } => {
accessor_name(value.node).map(|value| format!("{value}.{name}"))
}
_ => None,
}
}