/
fixes.rs
132 lines (118 loc) · 4.11 KB
/
fixes.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
use libcst_native::{
Codegen, CodegenState, Expression, ParenthesizableWhitespace, SmallStatement, Statement,
};
use rustpython_ast::{Expr, Keyword, Location};
use rustpython_parser::lexer;
use rustpython_parser::lexer::Tok;
use crate::ast::types::Range;
use crate::autofix::Fix;
use crate::source_code_locator::SourceCodeLocator;
/// Generate a fix to remove a base from a `ClassDef` statement.
pub fn remove_class_def_base(
locator: &SourceCodeLocator,
stmt_at: Location,
expr_at: Location,
bases: &[Expr],
keywords: &[Keyword],
) -> Option<Fix> {
let contents = locator.slice_source_code_at(&stmt_at);
// Case 1: `object` is the only base.
if bases.len() == 1 && keywords.is_empty() {
let mut fix_start = None;
let mut fix_end = None;
let mut count: usize = 0;
for (start, tok, end) in lexer::make_tokenizer_located(&contents, stmt_at).flatten() {
if matches!(tok, Tok::Lpar) {
if count == 0 {
fix_start = Some(start);
}
count += 1;
}
if matches!(tok, Tok::Rpar) {
count -= 1;
if count == 0 {
fix_end = Some(end);
break;
}
}
}
return match (fix_start, fix_end) {
(Some(start), Some(end)) => Some(Fix::deletion(start, end)),
_ => None,
};
}
if bases
.iter()
.map(|node| node.location)
.chain(keywords.iter().map(|node| node.location))
.any(|location| location > expr_at)
{
// Case 2: `object` is _not_ the last node.
let mut fix_start: Option<Location> = None;
let mut fix_end: Option<Location> = None;
let mut seen_comma = false;
for (start, tok, end) in lexer::make_tokenizer_located(&contents, stmt_at).flatten() {
if seen_comma {
if matches!(tok, Tok::Newline) {
fix_end = Some(end);
} else {
fix_end = Some(start);
}
break;
}
if start == expr_at {
fix_start = Some(start);
}
if fix_start.is_some() && matches!(tok, Tok::Comma) {
seen_comma = true;
}
}
match (fix_start, fix_end) {
(Some(start), Some(end)) => Some(Fix::replacement(String::new(), start, end)),
_ => None,
}
} else {
// Case 3: `object` is the last node, so we have to find the last token that
// isn't a comma.
let mut fix_start: Option<Location> = None;
let mut fix_end: Option<Location> = None;
for (start, tok, end) in lexer::make_tokenizer_located(&contents, stmt_at).flatten() {
if start == expr_at {
fix_end = Some(end);
break;
}
if matches!(tok, Tok::Comma) {
fix_start = Some(start);
}
}
match (fix_start, fix_end) {
(Some(start), Some(end)) => Some(Fix::replacement(String::new(), start, end)),
_ => None,
}
}
}
/// Generate a fix to remove arguments from a `super` call.
pub fn remove_super_arguments(locator: &SourceCodeLocator, expr: &Expr) -> Option<Fix> {
let range = Range::from_located(expr);
let contents = locator.slice_source_code_range(&range);
let mut tree = libcst_native::parse_module(&contents, None).ok()?;
let Statement::Simple(body) = tree.body.first_mut()? else {
return None;
};
let SmallStatement::Expr(body) = body.body.first_mut()? else {
return None;
};
let Expression::Call(body) = &mut body.value else {
return None;
};
body.args = vec![];
body.whitespace_before_args = ParenthesizableWhitespace::default();
body.whitespace_after_func = ParenthesizableWhitespace::default();
let mut state = CodegenState::default();
tree.codegen(&mut state);
Some(Fix::replacement(
state.to_string(),
range.location,
range.end_location,
))
}