Handle parentheses when formatting slice expressions (#5882)

**Summary** Fix the formatter crash with `x[(1) :: ]` and related code.

**Problem** For assigning comments in slices in subscripts, we need to
find the positions of the colons to assign comments before and after the
colon to the respective lower/upper/step node (or dangling in that
section). Formatting `x[(1) :: ]` was broken because we were looking for
a `:` after the `1` but didn't consider that there could be a `)`
outside the range of the lower node, which contains just the `1` and no
optional parentheses.

**Solution** Use the simple tokenizer directly and skip all closing
parentheses.

**Test Plan** I added regression tests.

Closes #5733
This commit is contained in:
konsti 2023-07-19 17:25:25 +02:00 committed by GitHub
parent 63ed7a31e8
commit a51606a10a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 31 additions and 14 deletions

View file

@ -5,7 +5,7 @@ use rustpython_parser::ast::{Expr, Ranged};
use ruff_formatter::prelude::{hard_line_break, line_suffix_boundary, space, text};
use ruff_formatter::{write, Buffer, Format, FormatError, FormatResult};
use ruff_python_ast::node::{AnyNodeRef, AstNode};
use ruff_python_whitespace::{first_non_trivia_token, Token, TokenKind};
use ruff_python_whitespace::{SimpleTokenizer, Token, TokenKind};
use crate::comments::{dangling_comments, SourceComment};
use crate::context::PyFormatContext;
@ -162,9 +162,12 @@ pub(crate) fn find_colons(
let after_lower = lower
.as_ref()
.map_or(range.start(), |lower| lower.range().end());
let first_colon = first_non_trivia_token(after_lower, contents).ok_or(
FormatError::syntax_error("Din't find any token for slice first colon"),
)?;
let mut tokens = SimpleTokenizer::new(contents, TextRange::new(after_lower, range.end()))
.skip_trivia()
.skip_while(|token| token.kind == TokenKind::RParen);
let first_colon = tokens.next().ok_or(FormatError::syntax_error(
"Din't find any token for slice first colon",
))?;
if first_colon.kind != TokenKind::Colon {
return Err(FormatError::syntax_error(
"slice first colon token was not a colon",
@ -174,16 +177,16 @@ pub(crate) fn find_colons(
let after_upper = upper
.as_ref()
.map_or(first_colon.end(), |upper| upper.range().end());
// At least the closing bracket must exist, so there must be a token there
let next_token = first_non_trivia_token(after_upper, contents).ok_or(
FormatError::syntax_error("Din't find any token for slice second colon"),
)?;
let second_colon = if next_token.kind == TokenKind::Colon {
debug_assert!(
next_token.range.start() < range.end(),
"The next token in a slice must either be a colon or the closing bracket"
);
Some(next_token)
let mut tokens = SimpleTokenizer::new(contents, TextRange::new(after_upper, range.end()))
.skip_trivia()
.skip_while(|token| token.kind == TokenKind::RParen);
let second_colon = if let Some(token) = tokens.next() {
if token.kind != TokenKind::Colon {
return Err(FormatError::syntax_error(
"Expected a colon for the second colon token",
));
}
Some(token)
} else {
None
};