Update the formatter code for the new AST design

astral-sh · Nov 22, 2023 · bcabcca · bcabcca
1 parent 6d6e1c7
commit bcabcca
Show file tree

Hide file tree

Showing 11 changed files with 134 additions and 233 deletions.
diff --git a/crates/ruff_python_ast/src/expression.rs b/crates/ruff_python_ast/src/expression.rs
@@ -362,6 +362,31 @@ impl Ranged for LiteralExpressionRef<'_> {
     }
 }
 
+impl<'a> From<LiteralExpressionRef<'a>> for AnyNodeRef<'a> {
+    fn from(value: LiteralExpressionRef<'a>) -> Self {
+        match value {
+            LiteralExpressionRef::StringLiteral(expression) => {
+                AnyNodeRef::ExprStringLiteral(expression)
+            }
+            LiteralExpressionRef::BytesLiteral(expression) => {
+                AnyNodeRef::ExprBytesLiteral(expression)
+            }
+            LiteralExpressionRef::NumberLiteral(expression) => {
+                AnyNodeRef::ExprNumberLiteral(expression)
+            }
+            LiteralExpressionRef::BooleanLiteral(expression) => {
+                AnyNodeRef::ExprBooleanLiteral(expression)
+            }
+            LiteralExpressionRef::NoneLiteral(expression) => {
+                AnyNodeRef::ExprNoneLiteral(expression)
+            }
+            LiteralExpressionRef::EllipsisLiteral(expression) => {
+                AnyNodeRef::ExprEllipsisLiteral(expression)
+            }
+        }
+    }
+}
+
 impl LiteralExpressionRef<'_> {
     /// Returns `true` if the literal is either a string or bytes literal that
     /// is implicitly concatenated.

diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs
@@ -283,13 +283,13 @@ fn handle_enclosed_comment<'a>(
         AnyNodeRef::StmtWith(with_) => handle_with_comment(comment, with_),
         AnyNodeRef::ExprCall(_) => handle_call_comment(comment),
         AnyNodeRef::ExprStringLiteral(_) => {
-            if let Some(AnyNodeRef::ExprFString(fstring)) = comment.enclosing_parent() {
+            if let Some(AnyNodeRef::FString(fstring)) = comment.enclosing_parent() {
                 CommentPlacement::dangling(fstring, comment)
             } else {
                 CommentPlacement::Default(comment)
             }
         }
-        AnyNodeRef::ExprFString(fstring) => CommentPlacement::dangling(fstring, comment),
+        AnyNodeRef::FString(fstring) => CommentPlacement::dangling(fstring, comment),
         AnyNodeRef::ExprList(_)
         | AnyNodeRef::ExprSet(_)
         | AnyNodeRef::ExprListComp(_)

diff --git a/crates/ruff_python_formatter/src/expression/expr_bin_op.rs b/crates/ruff_python_formatter/src/expression/expr_bin_op.rs
@@ -35,13 +35,13 @@ impl NeedsParentheses for ExprBinOp {
     ) -> OptionalParentheses {
         if parent.is_expr_await() {
             OptionalParentheses::Always
-        } else if self.left.is_literal_expr() {
+        } else if let Some(literal_expr) = self.left.as_literal_expr() {
             // Multiline strings are guaranteed to never fit, avoid adding unnecessary parentheses
-            if !self.left.is_implicit_concatenated_string()
-                && is_multiline_string(self.left.as_ref().into(), context.source())
+            if !literal_expr.is_implicit_concatenated()
+                && is_multiline_string(literal_expr.into(), context.source())
                 && has_parentheses(&self.right, context).is_some()
                 && !context.comments().has_dangling(self)
-                && !context.comments().has(self.left.as_ref())
+                && !context.comments().has(literal_expr)
                 && !context.comments().has(self.right.as_ref())
             {
                 OptionalParentheses::Never

diff --git a/crates/ruff_python_formatter/src/expression/expr_bytes_literal.rs b/crates/ruff_python_formatter/src/expression/expr_bytes_literal.rs
@@ -31,7 +31,7 @@ impl NeedsParentheses for ExprBytesLiteral {
         _parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
-        if self.implicit_concatenated {
+        if self.value.is_implicit_concatenated() {
             OptionalParentheses::Multiline
         } else if is_multiline_string(self.into(), context.source()) {
             OptionalParentheses::Never

diff --git a/crates/ruff_python_formatter/src/expression/expr_compare.rs b/crates/ruff_python_formatter/src/expression/expr_compare.rs
@@ -37,11 +37,11 @@ impl NeedsParentheses for ExprCompare {
     ) -> OptionalParentheses {
         if parent.is_expr_await() {
             OptionalParentheses::Always
-        } else if self.left.is_literal_expr() {
+        } else if let Some(literal_expr) = self.left.as_literal_expr() {
             // Multiline strings are guaranteed to never fit, avoid adding unnecessary parentheses
-            if !self.left.is_implicit_concatenated_string()
-                && is_multiline_string(self.left.as_ref().into(), context.source())
-                && !context.comments().has(self.left.as_ref())
+            if !literal_expr.is_implicit_concatenated()
+                && is_multiline_string(literal_expr.into(), context.source())
+                && !context.comments().has(literal_expr)
                 && self.comparators.first().is_some_and(|right| {
                     has_parentheses(right, context).is_some() && !context.comments().has(right)
                 })

diff --git a/crates/ruff_python_formatter/src/expression/expr_f_string.rs b/crates/ruff_python_formatter/src/expression/expr_f_string.rs
@@ -34,7 +34,7 @@ impl NeedsParentheses for ExprFString {
         _parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
-        if self.implicit_concatenated {
+        if self.value.is_implicit_concatenated() {
             OptionalParentheses::Multiline
         } else if memchr2(b'\n', b'\r', context.source()[self.range].as_bytes()).is_none() {
             OptionalParentheses::BestFit

diff --git a/crates/ruff_python_formatter/src/expression/expr_string_literal.rs b/crates/ruff_python_formatter/src/expression/expr_string_literal.rs
@@ -46,7 +46,7 @@ impl NeedsParentheses for ExprStringLiteral {
         _parent: AnyNodeRef,
         context: &PyFormatContext,
     ) -> OptionalParentheses {
-        if self.implicit_concatenated {
+        if self.value.is_implicit_concatenated() {
             OptionalParentheses::Multiline
         } else if is_multiline_string(self.into(), context.source()) {
             OptionalParentheses::Never

diff --git a/crates/ruff_python_formatter/src/expression/mod.rs b/crates/ruff_python_formatter/src/expression/mod.rs
@@ -804,18 +804,17 @@ impl<'input> CanOmitOptionalParenthesesVisitor<'input> {
                 return;
             }
 
-            Expr::StringLiteral(ast::ExprStringLiteral {
-                implicit_concatenated: true,
-                ..
-            })
-            | Expr::BytesLiteral(ast::ExprBytesLiteral {
-                implicit_concatenated: true,
-                ..
-            })
-            | Expr::FString(ast::ExprFString {
-                implicit_concatenated: true,
-                ..
-            }) => {
+            Expr::StringLiteral(ast::ExprStringLiteral { value, .. })
+                if value.is_implicit_concatenated() =>
+            {
+                self.update_max_precedence(OperatorPrecedence::String);
+            }
+            Expr::BytesLiteral(ast::ExprBytesLiteral { value, .. })
+                if value.is_implicit_concatenated() =>
+            {
+                self.update_max_precedence(OperatorPrecedence::String);
+            }
+            Expr::FString(ast::ExprFString { value, .. }) if value.is_implicit_concatenated() => {
                 self.update_max_precedence(OperatorPrecedence::String);
             }