Break

astral-sh · Aug 23, 2023 · 635c173 · 635c173
1 parent e3ce944
commit 635c173
Show file tree

Hide file tree

Showing 12 changed files with 159 additions and 88 deletions.
diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs
@@ -211,7 +211,6 @@ fn handle_enclosed_comment<'a>(
             handle_leading_class_with_decorators_comment(comment, class_def)
         }
         AnyNodeRef::StmtImportFrom(import_from) => handle_import_from_comment(comment, import_from),
-        AnyNodeRef::MatchCase(match_case) => handle_match_case_comment(comment, match_case),
         AnyNodeRef::StmtWith(with_) => handle_with_comment(comment, with_),
         AnyNodeRef::ExprConstant(_) => {
             if let Some(AnyNodeRef::ExprFString(fstring)) = comment.enclosing_parent() {
@@ -1346,28 +1345,6 @@ fn handle_import_from_comment<'a>(
     }
 }
 
-/// Attach an enclosed end-of-line comment to a [`MatchCase`].
-///
-/// For example, given:
-/// ```python
-/// case (  # comment
-///    pattern
-/// ):
-///     ...
-/// ```
-///
-/// The comment will be attached to the [`MatchCase`] node as a dangling comment.
-fn handle_match_case_comment<'a>(
-    comment: DecoratedComment<'a>,
-    match_case: &'a MatchCase,
-) -> CommentPlacement<'a> {
-    if comment.line_position().is_end_of_line() && comment.start() < match_case.pattern.start() {
-        CommentPlacement::dangling(comment.enclosing_node(), comment)
-    } else {
-        CommentPlacement::Default(comment)
-    }
-}
-
 /// Attach an enclosed end-of-line comment to a [`ast::StmtWith`].
 ///
 /// For example, given:

diff --git a/crates/ruff_python_formatter/src/other/match_case.rs b/crates/ruff_python_formatter/src/other/match_case.rs
@@ -1,13 +1,13 @@
 use ruff_formatter::{write, Buffer, FormatResult};
-use ruff_python_ast::{MatchCase, Pattern, Ranged};
-use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
-use ruff_text_size::TextRange;
+use ruff_python_ast::node::AstNode;
+use ruff_python_ast::MatchCase;
 
+use crate::builders::parenthesize_if_expands;
 use crate::comments::SourceComment;
-use crate::expression::parentheses::{parenthesized, Parentheses};
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parentheses};
 use crate::prelude::*;
 use crate::statement::clause::{clause_body, clause_header, ClauseHeader};
-use crate::{FormatError, FormatNodeRule, PyFormatter};
+use crate::{FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
 pub struct FormatMatchCase;
@@ -21,13 +21,8 @@ impl FormatNodeRule<MatchCase> for FormatMatchCase {
             body,
         } = item;
 
-        // Distinguish dangling comments that appear on the open parenthesis from those that
-        // appear on the trailing colon.
         let comments = f.context().comments().clone();
         let dangling_item_comments = comments.dangling(item);
-        let (open_parenthesis_comments, trailing_colon_comments) = dangling_item_comments.split_at(
-            dangling_item_comments.partition_point(|comment| comment.start() < pattern.start()),
-        );
 
         write!(
             f,
@@ -38,16 +33,26 @@ impl FormatNodeRule<MatchCase> for FormatMatchCase {
                     &format_with(|f| {
                         write!(f, [text("case"), space()])?;
 
-                        if is_match_case_pattern_parenthesized(item, pattern, f.context())? {
-                            parenthesized(
-                                "(",
-                                &pattern.format().with_options(Parentheses::Never),
-                                ")",
-                            )
-                            .with_dangling_comments(open_parenthesis_comments)
-                            .fmt(f)?;
+                        let has_comments = comments.has_leading(pattern)
+                            || comments.has_trailing_own_line(pattern);
+
+                        if has_comments {
+                            pattern.format().with_options(Parentheses::Always).fmt(f)?;
                         } else {
-                            pattern.format().fmt(f)?;
+                            match pattern.needs_parentheses(item.as_any_node_ref(), f.context()) {
+                                OptionalParentheses::Multiline => {
+                                    parenthesize_if_expands(
+                                        &pattern.format().with_options(Parentheses::Never),
+                                    )
+                                    .fmt(f)?;
+                                }
+                                OptionalParentheses::Always => {
+                                    pattern.format().with_options(Parentheses::Always).fmt(f)?;
+                                }
+                                OptionalParentheses::Never => {
+                                    pattern.format().with_options(Parentheses::Never).fmt(f)?;
+                                }
+                            }
                         }
 
                         if let Some(guard) = guard {
@@ -57,7 +62,7 @@ impl FormatNodeRule<MatchCase> for FormatMatchCase {
                         Ok(())
                     }),
                 ),
-                clause_body(body, trailing_colon_comments),
+                clause_body(body, dangling_item_comments),
             ]
         )
     }
@@ -71,33 +76,3 @@ impl FormatNodeRule<MatchCase> for FormatMatchCase {
         Ok(())
     }
 }
-
-fn is_match_case_pattern_parenthesized(
-    case: &MatchCase,
-    pattern: &Pattern,
-    context: &PyFormatContext,
-) -> FormatResult<bool> {
-    let mut tokenizer = SimpleTokenizer::new(
-        context.source(),
-        TextRange::new(case.start(), pattern.start()),
-    )
-    .skip_trivia();
-
-    let case_keyword = tokenizer.next().ok_or(FormatError::syntax_error(
-        "Expected a `case` keyword, didn't find any token",
-    ))?;
-
-    debug_assert_eq!(
-        case_keyword.kind(),
-        SimpleTokenKind::Case,
-        "Expected `case` keyword but at {case_keyword:?}"
-    );
-
-    match tokenizer.next() {
-        Some(left_paren) => {
-            debug_assert_eq!(left_paren.kind(), SimpleTokenKind::LParen);
-            Ok(true)
-        }
-        None => Ok(false),
-    }
-}
diff --git a/crates/ruff_python_formatter/src/pattern/mod.rs b/crates/ruff_python_formatter/src/pattern/mod.rs
@@ -1,8 +1,11 @@
 use ruff_formatter::{FormatOwnedWithRule, FormatRefWithRule, FormatRule, FormatRuleWithOptions};
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{Pattern, Ranged};
 use ruff_python_trivia::{first_non_trivia_token, SimpleToken, SimpleTokenKind, SimpleTokenizer};
 
-use crate::expression::parentheses::{parenthesized, Parentheses};
+use crate::expression::parentheses::{
+    parenthesized, NeedsParentheses, OptionalParentheses, Parentheses,
+};
 use crate::prelude::*;
 
 pub(crate) mod pattern_match_as;
@@ -115,3 +118,22 @@ fn is_pattern_parenthesized(pattern: &Pattern, contents: &str) -> bool {
         false
     }
 }
+
+impl NeedsParentheses for Pattern {
+    fn needs_parentheses(
+        &self,
+        parent: AnyNodeRef,
+        context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        match self {
+            Pattern::MatchValue(pattern) => pattern.needs_parentheses(parent, context),
+            Pattern::MatchSingleton(pattern) => pattern.needs_parentheses(parent, context),
+            Pattern::MatchSequence(pattern) => pattern.needs_parentheses(parent, context),
+            Pattern::MatchMapping(pattern) => pattern.needs_parentheses(parent, context),
+            Pattern::MatchClass(pattern) => pattern.needs_parentheses(parent, context),
+            Pattern::MatchStar(pattern) => pattern.needs_parentheses(parent, context),
+            Pattern::MatchAs(pattern) => pattern.needs_parentheses(parent, context),
+            Pattern::MatchOr(pattern) => pattern.needs_parentheses(parent, context),
+        }
+    }
+}
diff --git a/crates/ruff_python_formatter/src/pattern/pattern_match_as.rs b/crates/ruff_python_formatter/src/pattern/pattern_match_as.rs
@@ -1,7 +1,9 @@
 use ruff_formatter::{write, Buffer, FormatResult};
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::PatternMatchAs;
 
 use crate::comments::{dangling_comments, SourceComment};
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::prelude::*;
 use crate::{FormatNodeRule, PyFormatter};
 
@@ -56,3 +58,13 @@ impl FormatNodeRule<PatternMatchAs> for FormatPatternMatchAs {
         Ok(())
     }
 }
+
+impl NeedsParentheses for PatternMatchAs {
+    fn needs_parentheses(
+        &self,
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Multiline
+    }
+}
diff --git a/crates/ruff_python_formatter/src/pattern/pattern_match_class.rs b/crates/ruff_python_formatter/src/pattern/pattern_match_class.rs
@@ -1,6 +1,9 @@
 use ruff_formatter::{write, Buffer, FormatResult};
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::PatternMatchClass;
 
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
+use crate::prelude::*;
 use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
@@ -17,3 +20,13 @@ impl FormatNodeRule<PatternMatchClass> for FormatPatternMatchClass {
         )
     }
 }
+
+impl NeedsParentheses for PatternMatchClass {
+    fn needs_parentheses(
+        &self,
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Never
+    }
+}
diff --git a/crates/ruff_python_formatter/src/pattern/pattern_match_mapping.rs b/crates/ruff_python_formatter/src/pattern/pattern_match_mapping.rs
@@ -1,6 +1,9 @@
 use ruff_formatter::{write, Buffer, FormatResult};
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::PatternMatchMapping;
 
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
+use crate::prelude::*;
 use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
@@ -17,3 +20,13 @@ impl FormatNodeRule<PatternMatchMapping> for FormatPatternMatchMapping {
         )
     }
 }
+
+impl NeedsParentheses for PatternMatchMapping {
+    fn needs_parentheses(
+        &self,
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Never
+    }
+}
diff --git a/crates/ruff_python_formatter/src/pattern/pattern_match_or.rs b/crates/ruff_python_formatter/src/pattern/pattern_match_or.rs
@@ -1,6 +1,9 @@
 use ruff_formatter::{write, Buffer, FormatResult};
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::PatternMatchOr;
 
+use crate::context::PyFormatContext;
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
@@ -17,3 +20,13 @@ impl FormatNodeRule<PatternMatchOr> for FormatPatternMatchOr {
         )
     }
 }
+
+impl NeedsParentheses for PatternMatchOr {
+    fn needs_parentheses(
+        &self,
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Multiline
+    }
+}
diff --git a/crates/ruff_python_formatter/src/pattern/pattern_match_sequence.rs b/crates/ruff_python_formatter/src/pattern/pattern_match_sequence.rs
@@ -1,9 +1,13 @@
 use ruff_formatter::prelude::format_with;
 use ruff_formatter::{Format, FormatResult};
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::PatternMatchSequence;
 
 use crate::builders::PyFormatterExtensions;
-use crate::expression::parentheses::{empty_parenthesized, optional_parentheses, parenthesized};
+use crate::context::PyFormatContext;
+use crate::expression::parentheses::{
+    empty_parenthesized, optional_parentheses, parenthesized, NeedsParentheses, OptionalParentheses,
+};
 use crate::{FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
@@ -51,3 +55,13 @@ impl FormatNodeRule<PatternMatchSequence> for FormatPatternMatchSequence {
         }
     }
 }
+
+impl NeedsParentheses for PatternMatchSequence {
+    fn needs_parentheses(
+        &self,
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Never
+    }
+}
diff --git a/crates/ruff_python_formatter/src/pattern/pattern_match_singleton.rs b/crates/ruff_python_formatter/src/pattern/pattern_match_singleton.rs
@@ -1,6 +1,8 @@
 use crate::prelude::*;
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::{Constant, PatternMatchSingleton};
 
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::{FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
@@ -16,3 +18,13 @@ impl FormatNodeRule<PatternMatchSingleton> for FormatPatternMatchSingleton {
         }
     }
 }
+
+impl NeedsParentheses for PatternMatchSingleton {
+    fn needs_parentheses(
+        &self,
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Never
+    }
+}
diff --git a/crates/ruff_python_formatter/src/pattern/pattern_match_star.rs b/crates/ruff_python_formatter/src/pattern/pattern_match_star.rs
@@ -1,6 +1,9 @@
 use ruff_formatter::{write, Buffer, FormatResult};
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::PatternMatchStar;
 
+use crate::context::PyFormatContext;
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses};
 use crate::{not_yet_implemented_custom_text, FormatNodeRule, PyFormatter};
 
 #[derive(Default)]
@@ -17,3 +20,13 @@ impl FormatNodeRule<PatternMatchStar> for FormatPatternMatchStar {
         )
     }
 }
+
+impl NeedsParentheses for PatternMatchStar {
+    fn needs_parentheses(
+        &self,
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Never
+    }
+}
diff --git a/crates/ruff_python_formatter/src/pattern/pattern_match_value.rs b/crates/ruff_python_formatter/src/pattern/pattern_match_value.rs
@@ -1,6 +1,7 @@
+use ruff_python_ast::node::AnyNodeRef;
 use ruff_python_ast::PatternMatchValue;
 
-use crate::expression::parentheses::Parentheses;
+use crate::expression::parentheses::{NeedsParentheses, OptionalParentheses, Parentheses};
 use crate::prelude::*;
 use crate::{AsFormat, FormatNodeRule, PyFormatter};
 
@@ -13,3 +14,13 @@ impl FormatNodeRule<PatternMatchValue> for FormatPatternMatchValue {
         value.format().with_options(Parentheses::Never).fmt(f)
     }
 }
+
+impl NeedsParentheses for PatternMatchValue {
+    fn needs_parentheses(
+        &self,
+        _parent: AnyNodeRef,
+        _context: &PyFormatContext,
+    ) -> OptionalParentheses {
+        OptionalParentheses::Never
+    }
+}