formatter: WithNodeLevel helper (#6212)

This commit is contained in:
Micha Reiser 2023-07-31 23:22:17 +02:00 committed by GitHub
parent 615337a54d
commit 38b5726948
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 208 additions and 212 deletions

View file

@ -3,12 +3,14 @@ use std::cmp::Ordering;
use ruff_python_ast as ast;
use ruff_python_ast::{Expr, Operator};
use ruff_formatter::{FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions};
use ruff_formatter::{
write, FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions,
};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::visitor::preorder::{walk_expr, PreorderVisitor};
use crate::builders::parenthesize_if_expands;
use crate::context::NodeLevel;
use crate::context::{NodeLevel, WithNodeLevel};
use crate::expression::parentheses::{
is_expression_parenthesized, optional_parentheses, parenthesized, NeedsParentheses,
OptionalParentheses, Parentheses, Parenthesize,
@ -106,21 +108,16 @@ impl FormatRule<Expr, PyFormatContext<'_>> for FormatExpr {
if parenthesize {
parenthesized("(", &format_expr, ")").fmt(f)
} else {
let saved_level = match f.context().node_level() {
saved_level @ (NodeLevel::TopLevel | NodeLevel::CompoundStatement) => {
f.context_mut().set_node_level(NodeLevel::Expression(None));
Some(saved_level)
let level = match f.context().node_level() {
NodeLevel::TopLevel | NodeLevel::CompoundStatement => NodeLevel::Expression(None),
saved_level @ (NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression) => {
saved_level
}
NodeLevel::Expression(_) | NodeLevel::ParenthesizedExpression => None,
};
let result = Format::fmt(&format_expr, f);
let mut f = WithNodeLevel::new(level, f);
if let Some(saved_level) = saved_level {
f.context_mut().set_node_level(saved_level);
}
result
write!(f, [format_expr])
}
}
}