From a51606a10a6126d6fc074d2f85a8a4ee1c3e33a2 Mon Sep 17 00:00:00 2001
From: konsti <konstin@mailbox.org>
Date: Wed, 19 Jul 2023 17:25:25 +0200
Subject: [PATCH] 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
---
 .../test/fixtures/ruff/expression/slice.py    |  5 +++
 .../src/expression/expr_slice.rs              | 31 ++++++++++---------
 .../format@expression__slice.py.snap          |  9 ++++++
 3 files changed, 31 insertions(+), 14 deletions(-)

diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/slice.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/slice.py
index 8e0b32b72e..bbdfe19010 100644
--- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/slice.py
+++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/slice.py
@@ -83,3 +83,8 @@ e210 = "e"[a() : 1 :]
 
 # Regression test for https://github.com/astral-sh/ruff/issues/5605
 f = "f"[:,]
+
+# Regression test for https://github.com/astral-sh/ruff/issues/5733
+g1 = "g"[(1):(2)]
+g2 = "g"[(1):(2):(3)]
+
diff --git a/crates/ruff_python_formatter/src/expression/expr_slice.rs b/crates/ruff_python_formatter/src/expression/expr_slice.rs
index be5f8396b7..78f48a2e71 100644
--- a/crates/ruff_python_formatter/src/expression/expr_slice.rs
+++ b/crates/ruff_python_formatter/src/expression/expr_slice.rs
@@ -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
     };
diff --git a/crates/ruff_python_formatter/tests/snapshots/format@expression__slice.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@expression__slice.py.snap
index 5c9251a69d..a50771e6fb 100644
--- a/crates/ruff_python_formatter/tests/snapshots/format@expression__slice.py.snap
+++ b/crates/ruff_python_formatter/tests/snapshots/format@expression__slice.py.snap
@@ -89,6 +89,11 @@ e210 = "e"[a() : 1 :]
 
 # Regression test for https://github.com/astral-sh/ruff/issues/5605
 f = "f"[:,]
+
+# Regression test for https://github.com/astral-sh/ruff/issues/5733
+g1 = "g"[(1):(2)]
+g2 = "g"[(1):(2):(3)]
+
 ```
 
 ## Output
@@ -176,6 +181,10 @@ e210 = "e"[a() : 1 :]
 
 # Regression test for https://github.com/astral-sh/ruff/issues/5605
 f = "f"[:,]
+
+# Regression test for https://github.com/astral-sh/ruff/issues/5733
+g1 = "g"[(1):(2)]
+g2 = "g"[(1):(2):(3)]
 ```