Handle bracketed comments on sequence patterns (#6801)

## Summary

This PR ensures that we handle bracketed comments on sequences, like `#
comment` here:

```python
match x:
    case [ # comment
        1, 2
    ]:
        pass
```

The handling is very similar to other, similar nodes, except that we do
need some special logic to determine whether the sequence is
parenthesized, similar to our logic for tuples.

## Test Plan

`cargo test`
This commit is contained in:
Charlie Marsh 2023-08-25 00:03:27 -04:00 committed by GitHub
parent 474e8fbcd4
commit f754ad5898
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 93 additions and 32 deletions

View file

@ -15,6 +15,7 @@ use crate::expression::expr_tuple::is_tuple_parenthesized;
use crate::other::parameters::{
assign_argument_separator_comment_placement, find_parameter_separators,
};
use crate::pattern::pattern_match_sequence::SequenceType;
/// Manually attach comments to nodes that the default placement gets wrong.
pub(super) fn place_comment<'a>(
@ -179,6 +180,15 @@ fn handle_enclosed_comment<'a>(
AnyNodeRef::Comprehension(comprehension) => {
handle_comprehension_comment(comment, comprehension, locator)
}
AnyNodeRef::PatternMatchSequence(pattern_match_sequence) => {
if SequenceType::from_pattern(pattern_match_sequence, locator.contents())
.is_parenthesized()
{
handle_bracketed_end_of_line_comment(comment, locator)
} else {
CommentPlacement::Default(comment)
}
}
AnyNodeRef::ExprAttribute(attribute) => {
handle_attribute_comment(comment, attribute, locator)
}

View file

@ -1,7 +1,9 @@
use ruff_formatter::prelude::format_with;
use ruff_formatter::{Format, FormatResult};
use ruff_python_ast::node::AnyNodeRef;
use ruff_python_ast::PatternMatchSequence;
use ruff_python_ast::{PatternMatchSequence, Ranged};
use ruff_python_trivia::{SimpleTokenKind, SimpleTokenizer};
use ruff_text_size::TextRange;
use crate::builders::PyFormatterExtensions;
use crate::context::PyFormatContext;
@ -13,29 +15,19 @@ use crate::{FormatNodeRule, PyFormatter};
#[derive(Default)]
pub struct FormatPatternMatchSequence;
#[derive(Debug)]
enum SequenceType {
Tuple,
TupleNoParens,
List,
}
impl FormatNodeRule<PatternMatchSequence> for FormatPatternMatchSequence {
fn fmt_fields(&self, item: &PatternMatchSequence, f: &mut PyFormatter) -> FormatResult<()> {
let PatternMatchSequence { patterns, range } = item;
let sequence_type = match &f.context().source()[*range].chars().next() {
Some('(') => SequenceType::Tuple,
Some('[') => SequenceType::List,
_ => SequenceType::TupleNoParens,
};
let comments = f.context().comments().clone();
let dangling = comments.dangling(item);
let sequence_type = SequenceType::from_pattern(item, f.context().source());
if patterns.is_empty() {
return match sequence_type {
SequenceType::Tuple => empty_parenthesized("(", dangling, ")").fmt(f),
SequenceType::List => empty_parenthesized("[", dangling, "]").fmt(f),
SequenceType::TupleNoParens => {
unreachable!("If empty, it should be either tuple or list")
SequenceType::Tuple | SequenceType::TupleNoParens => {
empty_parenthesized("(", dangling, ")").fmt(f)
}
};
}
@ -65,3 +57,68 @@ impl NeedsParentheses for PatternMatchSequence {
OptionalParentheses::Never
}
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub(crate) enum SequenceType {
/// A list literal, e.g., `[1, 2, 3]`.
List,
/// A parenthesized tuple literal, e.g., `(1, 2, 3)`.
Tuple,
/// A tuple literal without parentheses, e.g., `1, 2, 3`.
TupleNoParens,
}
impl SequenceType {
pub(crate) fn from_pattern(pattern: &PatternMatchSequence, source: &str) -> SequenceType {
if source[pattern.range()].starts_with('[') {
SequenceType::List
} else if source[pattern.range()].starts_with('(') {
// If the pattern is empty, it must be a parenthesized tuple with no members. (This
// branch exists to differentiate between a tuple with and without its own parentheses,
// but a tuple without its own parentheses must have at least one member.)
let Some(elt) = pattern.patterns.first() else {
return SequenceType::Tuple;
};
// Count the number of open parentheses between the start of the pattern and the first
// element, and the number of close parentheses between the first element and its
// trailing comma. If the number of open parentheses is greater than the number of close
// parentheses,
// the pattern is parenthesized. For example, here, we have two parentheses before the
// first element, and one after it:
// ```python
// ((a), b, c)
// ```
//
// This algorithm successfully avoids false positives for cases like:
// ```python
// (a), b, c
// ```
let open_parentheses_count =
SimpleTokenizer::new(source, TextRange::new(pattern.start(), elt.start()))
.skip_trivia()
.filter(|token| token.kind() == SimpleTokenKind::LParen)
.count();
// Count the number of close parentheses.
let close_parentheses_count =
SimpleTokenizer::new(source, TextRange::new(elt.end(), elt.end()))
.skip_trivia()
.take_while(|token| token.kind() != SimpleTokenKind::Comma)
.filter(|token| token.kind() == SimpleTokenKind::RParen)
.count();
if open_parentheses_count > close_parentheses_count {
SequenceType::Tuple
} else {
SequenceType::TupleNoParens
}
} else {
SequenceType::TupleNoParens
}
}
pub(crate) fn is_parenthesized(self) -> bool {
matches!(self, SequenceType::List | SequenceType::Tuple)
}
}