mirror of https://github.com/astral-sh/ruff
32 lines
931 B
Rust
32 lines
931 B
Rust
use crate::core::visitor;
|
|
use crate::core::visitor::Visitor;
|
|
use crate::cst::{Expr, Stmt};
|
|
use crate::trivia::{decorate_trivia, TriviaIndex, TriviaToken};
|
|
|
|
struct AttachmentVisitor {
|
|
index: TriviaIndex,
|
|
}
|
|
|
|
impl<'a> Visitor<'a> for AttachmentVisitor {
|
|
fn visit_stmt(&mut self, stmt: &'a mut Stmt) {
|
|
let trivia = self.index.stmt.remove(&stmt.id());
|
|
if let Some(comments) = trivia {
|
|
stmt.trivia.extend(comments);
|
|
}
|
|
visitor::walk_stmt(self, stmt);
|
|
}
|
|
|
|
fn visit_expr(&mut self, expr: &'a mut Expr) {
|
|
let trivia = self.index.expr.remove(&expr.id());
|
|
if let Some(comments) = trivia {
|
|
expr.trivia.extend(comments);
|
|
}
|
|
visitor::walk_expr(self, expr);
|
|
}
|
|
}
|
|
|
|
pub fn attach(python_cst: &mut [Stmt], trivia: Vec<TriviaToken>) {
|
|
let index = decorate_trivia(trivia, python_cst);
|
|
AttachmentVisitor { index }.visit_body(python_cst);
|
|
}
|