Format Slice Expressions

crates/ruff_python_formatter/resources/test/fixtures/ruff/expression/slice.py

@@ -0,0 +1,82 @@
+# Handle comments both when lower and upper exist and when they don't
+a1 = "a"[
+    # a
+    1  # b
+    :  # c
+    2  # d
+]
+a2 = "a"[
+    # a
+      # b
+    :  # c
+      # d
+]
+
+# Check all places where comments can exist
+b1 = "b"[ # a
+    # b
+    1 # c
+    # d
+    : # e
+    # f
+    2 # g
+    # h
+    : # i
+    # j
+    3 # k
+    # l
+]
+
+# Handle the spacing from the colon correctly with upper leading comments
+c1 = "c"[
+    1
+    :  # e
+    # f
+    2
+]
+c2 = "c"[
+    1
+    :  # e
+    2
+]
+c3 = "c"[
+    1
+    :
+    # f
+    2
+]
+c4 = "c"[
+    1
+    : # f
+    2
+]
+
+# End of line comments
+d1 = "d"[ # comment
+    :
+]
+d2 = "d"[  # comment
+    1:
+]
+d3 = "d"[
+    1  # comment
+    :
+]
+
+# Spacing around the colon(s)
+def a():
+    ...
+
+e00 = "e"[:]
+e01 = "e"[:1]
+e02 = "e"[: a()]
+e10 = "e"[1:]
+e11 = "e"[1:1]
+e12 = "e"[1 : a()]
+e20 = "e"[a() :]
+e21 = "e"[a() : 1]
+e22 = "e"[a() : a()]
+e200 = "e"[a() :: ]
+e201 = "e"[a() :: 1]
+e202 = "e"[a() :: a()]
+e210 = "e"[a() : 1 :]

crates/ruff_python_formatter/src/comments/placement.rs

@@ -1,16 +1,14 @@
-use std::cmp::Ordering;
-
-use ruff_text_size::{TextRange, TextSize};
-use rustpython_parser::ast::Ranged;
-
-use ruff_python_ast::node::AnyNodeRef;
+use crate::comments::visitor::{CommentPlacement, DecoratedComment};
+use crate::comments::CommentLinePosition;
+use crate::expression::expr_slice::{assign_comment_in_slice, ExprSliceCommentSection};
+use crate::trivia::{first_non_trivia_token_rev, SimpleTokenizer, Token, TokenKind};
+use ruff_python_ast::node::{AnyNodeRef, AstNode};
 use ruff_python_ast::source_code::Locator;
 use ruff_python_ast::whitespace;
 use ruff_python_whitespace::{PythonWhitespace, UniversalNewlines};
-
-use crate::comments::visitor::{CommentPlacement, DecoratedComment};
-use crate::comments::CommentLinePosition;
-use crate::trivia::{SimpleTokenizer, Token, TokenKind};
+use ruff_text_size::{TextRange, TextSize};
+use rustpython_parser::ast::{Expr, ExprSlice, Ranged};
+use std::cmp::Ordering;
 
 /// Implements the custom comment placement logic.
 pub(super) fn place_comment<'a>(
@@ -32,6 +30,7 @@ pub(super) fn place_comment<'a>(
     .or_else(|comment| handle_positional_only_arguments_separator_comment(comment, locator))
     .or_else(|comment| handle_trailing_binary_expression_left_or_operator_comment(comment, locator))
     .or_else(handle_leading_function_with_decorators_comment)
+    .or_else(|comment| handle_slice_comments(comment, locator))
     .or_else(|comment| handle_dict_unpacking_comment(comment, locator))
 }
 
@@ -829,6 +828,87 @@ fn handle_module_level_own_line_comment_before_class_or_function_comment<'a>(
     }
 }
 
+/// Handles the attaching comments left or right of the colon in a slice as trailing comment of the
+/// preceding node or leading comment of the following node respectively.
+/// ```python
+/// a = "input"[
+///     1 # c
+///     # d
+///     :2
+/// ]
+/// ```
+fn handle_slice_comments<'a>(
+    comment: DecoratedComment<'a>,
+    locator: &Locator,
+) -> CommentPlacement<'a> {
+    let expr_slice = match comment.enclosing_node() {
+        AnyNodeRef::ExprSlice(expr_slice) => expr_slice,
+        AnyNodeRef::ExprSubscript(expr_subscript) => {
+            if expr_subscript.value.end() < expr_subscript.slice.start() {
+                if let Expr::Slice(expr_slice) = expr_subscript.slice.as_ref() {
+                    expr_slice
+                } else {
+                    return CommentPlacement::Default(comment);
+                }
+            } else {
+                return CommentPlacement::Default(comment);
+            }
+        }
+        _ => return CommentPlacement::Default(comment),
+    };
+
+    let ExprSlice {
+        range: _,
+        lower,
+        upper,
+        step,
+    } = expr_slice;
+
+    // Check for `foo[ # comment`, but only if they are on the same line
+    let after_lbracket = matches!(
+        first_non_trivia_token_rev(comment.slice().start(), locator.contents()),
+        Some(Token {
+            kind: TokenKind::LBracket,
+            ..
+        })
+    );
+    if comment.line_position().is_end_of_line() && after_lbracket {
+        // Keep comments after the opening bracket there by formatting them outside the
+        // soft block indent
+        // ```python
+        // "a"[ # comment
+        //     1:
+        // ]
+        // ```
+        debug_assert!(
+            matches!(comment.enclosing_node(), AnyNodeRef::ExprSubscript(_)),
+            "{:?}",
+            comment.enclosing_node()
+        );
+        return CommentPlacement::dangling(comment.enclosing_node(), comment);
+    }
+
+    let assignment =
+        assign_comment_in_slice(comment.slice().range(), locator.contents(), expr_slice);
+    let node = match assignment {
+        ExprSliceCommentSection::Lower => lower,
+        ExprSliceCommentSection::Upper => upper,
+        ExprSliceCommentSection::Step => step,
+    };
+
+    if let Some(node) = node {
+        if comment.slice().start() < node.start() {
+            CommentPlacement::leading(node.as_ref().into(), comment)
+        } else {
+            // If a trailing comment is an end of line comment that's fine because we have a node
+            // ahead of it
+            CommentPlacement::trailing(node.as_ref().into(), comment)
+        }
+    } else {
+        CommentPlacement::dangling(expr_slice.as_any_node_ref(), comment)
+    }
+}
+
 /// Finds the offset of the `/` that separates the positional only and arguments from the other arguments.
 /// Returns `None` if the positional only separator `/` isn't present in the specified range.
 fn find_pos_only_slash_offset(