formatter: WithNodeLevel helper (#6212)

crates/ruff_python_formatter/src/expression/expr_subscript.rs

@@ -4,8 +4,8 @@ use ruff_formatter::{format_args, write};
 use ruff_python_ast::node::{AnyNodeRef, AstNode};
 
 use crate::comments::trailing_comments;
-use crate::context::NodeLevel;
 use crate::context::PyFormatContext;
+use crate::context::{NodeLevel, WithNodeLevel};
 use crate::expression::expr_tuple::TupleParentheses;
 use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
@@ -30,34 +30,22 @@ impl FormatNodeRule<ExprSubscript> for FormatExprSubscript {
             "A subscript expression can only have a single dangling comment, the one after the bracket"
         );
 
-        if let NodeLevel::Expression(Some(group_id)) = f.context().node_level() {
+        if let NodeLevel::Expression(Some(_)) = f.context().node_level() {
             // Enforce the optional parentheses for parenthesized values.
-            f.context_mut().set_node_level(NodeLevel::Expression(None));
-            let result = value.format().fmt(f);
-            f.context_mut()
-                .set_node_level(NodeLevel::Expression(Some(group_id)));
-            result?;
+            let mut f = WithNodeLevel::new(NodeLevel::Expression(None), f);
+            write!(f, [value.format()])?;
         } else {
             value.format().fmt(f)?;
         }
 
         let format_slice = format_with(|f: &mut PyFormatter| {
-            let saved_level = f.context().node_level();
-            f.context_mut()
-                .set_node_level(NodeLevel::ParenthesizedExpression);
+            let mut f = WithNodeLevel::new(NodeLevel::ParenthesizedExpression, f);
 
-            let result = if let Expr::Tuple(tuple) = slice.as_ref() {
-                tuple
-                    .format()
-                    .with_options(TupleParentheses::Preserve)
-                    .fmt(f)
+            if let Expr::Tuple(tuple) = slice.as_ref() {
+                write!(f, [tuple.format().with_options(TupleParentheses::Preserve)])
             } else {
-                slice.format().fmt(f)
-            };
-
-            f.context_mut().set_node_level(saved_level);
-
-            result
+                write!(f, [slice.format()])
+            }
         });
 
         write!(

crates/ruff_python_formatter/src/expression/mod.rs

@@ -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])
         }
     }
 }

crates/ruff_python_formatter/src/expression/parentheses.rs

@@ -5,7 +5,7 @@ use ruff_formatter::{format_args, write, Argument, Arguments};
 use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_trivia::{first_non_trivia_token, SimpleToken, SimpleTokenKind, SimpleTokenizer};
 
-use crate::context::NodeLevel;
+use crate::context::{NodeLevel, WithNodeLevel};
 use crate::prelude::*;
 
 #[derive(Copy, Clone, Debug, Eq, PartialEq)]
@@ -134,23 +134,20 @@ impl<'ast> Format<PyFormatContext<'ast>> for FormatParenthesized<'_, 'ast> {
 
         let current_level = f.context().node_level();
 
-        f.context_mut()
-            .set_node_level(NodeLevel::ParenthesizedExpression);
+        let mut f = WithNodeLevel::new(NodeLevel::ParenthesizedExpression, f);
 
-        let result = if let NodeLevel::Expression(Some(group_id)) = current_level {
+        if let NodeLevel::Expression(Some(group_id)) = current_level {
             // Use fits expanded if there's an enclosing group that adds the optional parentheses.
             // This ensures that expanding this parenthesized expression does not expand the optional parentheses group.
-            fits_expanded(&inner)
-                .with_condition(Some(Condition::if_group_fits_on_line(group_id)))
-                .fmt(f)
+            write!(
+                f,
+                [fits_expanded(&inner)
+                    .with_condition(Some(Condition::if_group_fits_on_line(group_id)))]
+            )
         } else {
             // It's not necessary to wrap the content if it is not inside of an optional_parentheses group.
-            inner.fmt(f)
-        };
-
-        f.context_mut().set_node_level(current_level);
-
-        result
+            write!(f, [inner])
+        }
     }
 }
 
@@ -173,35 +170,30 @@ pub(crate) struct FormatOptionalParentheses<'content, 'ast> {
 
 impl<'ast> Format<PyFormatContext<'ast>> for FormatOptionalParentheses<'_, 'ast> {
     fn fmt(&self, f: &mut Formatter<PyFormatContext<'ast>>) -> FormatResult<()> {
-        let saved_level = f.context().node_level();
-
         // The group id is used as a condition in [`in_parentheses_only`] to create a conditional group
         // that is only active if the optional parentheses group expands.
         let parens_id = f.group_id("optional_parentheses");
 
-        f.context_mut()
-            .set_node_level(NodeLevel::Expression(Some(parens_id)));
+        let mut f = WithNodeLevel::new(NodeLevel::Expression(Some(parens_id)), f);
 
         // We can't use `soft_block_indent` here because that would always increment the indent,
         // even if the group does not break (the indent is not soft). This would result in
         // too deep indentations if a `parenthesized` group expands. Using `indent_if_group_breaks`
         // gives us the desired *soft* indentation that is only present if the optional parentheses
         // are shown.
-        let result = group(&format_args![
-            if_group_breaks(&text("(")),
-            indent_if_group_breaks(
-                &format_args![soft_line_break(), Arguments::from(&self.content)],
-                parens_id
-            ),
-            soft_line_break(),
-            if_group_breaks(&text(")"))
-        ])
-        .with_group_id(Some(parens_id))
-        .fmt(f);
-
-        f.context_mut().set_node_level(saved_level);
-
-        result
+        write!(
+            f,
+            [group(&format_args![
+                if_group_breaks(&text("(")),
+                indent_if_group_breaks(
+                    &format_args![soft_line_break(), Arguments::from(&self.content)],
+                    parens_id
+                ),
+                soft_line_break(),
+                if_group_breaks(&text(")"))
+            ])
+            .with_group_id(Some(parens_id))]
+        )
     }
 }