Add support for trailing colons in slice expressions (#3077)

crates/ruff_python_formatter/src/format/comprehension.rs

@@ -22,6 +22,10 @@ impl Format<ASTFormatContext<'_>> for FormatComprehension<'_> {
         let comprehension = self.item;
 
         write!(f, [soft_line_break_or_space()])?;
+        if comprehension.is_async > 0 {
+            write!(f, [text("async")])?;
+            write!(f, [space()])?;
+        }
         write!(f, [text("for")])?;
         write!(f, [space()])?;
         // TODO(charlie): If this is an unparenthesized tuple, we need to avoid expanding it.

crates/ruff_python_formatter/src/format/expr.rs

@@ -218,6 +218,17 @@ fn format_slice(
             write!(f, [space()])?;
         }
         write!(f, [step.format()])?;
+    } else {
+        let magic_trailing_colon = expr
+            .trivia
+            .iter()
+            .any(|c| matches!(c.kind, TriviaKind::MagicTrailingColon));
+        if magic_trailing_colon {
+            if !is_simple && upper.is_some() {
+                write!(f, [space()])?;
+            }
+            write!(f, [text(":")])?;
+        }
     }
 
     Ok(())

crates/ruff_python_formatter/src/lib.rs

@@ -70,6 +70,7 @@ mod tests {
     #[test_case(Path::new("simple_cases/beginning_backslash.py"); "beginning_backslash")]
     #[test_case(Path::new("simple_cases/import_spacing.py"); "import_spacing")]
     #[test_case(Path::new("simple_cases/power_op_spacing.py"); "power_op_spacing")]
+    #[test_case(Path::new("simple_cases/slices.py"); "slices")]
     fn passing(path: &Path) -> Result<()> {
         let snapshot = format!("{}", path.display());
         let content = std::fs::read_to_string(test_resource_path(

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__simple_cases__slices.py.snap

@@ -34,4 +34,3 @@ ham[lower:upper], ham[lower:upper:], ham[lower::step]
 # ham[lower+offset : upper+offset]
 ham[: upper_fn(x) : step_fn(x)], ham[:: step_fn(x)]
 ham[lower + offset : upper + offset]
-

crates/ruff_python_formatter/src/trivia.rs

@@ -31,6 +31,7 @@ pub enum TriviaTokenKind {
     OwnLineComment,
     EndOfLineComment,
     MagicTrailingComma,
+    MagicTrailingColon,
     EmptyLine,
     Parentheses,
 }
@@ -66,6 +67,7 @@ pub enum TriviaKind {
     /// ```
     EndOfLineComment(Range),
     MagicTrailingComma,
+    MagicTrailingColon,
     EmptyLine,
     Parentheses,
 }
@@ -100,6 +102,10 @@ impl Trivia {
                 kind: TriviaKind::MagicTrailingComma,
                 relationship,
             },
+            TriviaTokenKind::MagicTrailingColon => Self {
+                kind: TriviaKind::MagicTrailingColon,
+                relationship,
+            },
             TriviaTokenKind::EmptyLine => Self {
                 kind: TriviaKind::EmptyLine,
                 relationship,
@@ -163,6 +169,12 @@ pub fn extract_trivia_tokens(lxr: &[LexResult]) -> Vec<TriviaToken> {
                         end: *prev_end,
                         kind: TriviaTokenKind::MagicTrailingComma,
                     });
+                } else if prev_tok == &Tok::Colon {
+                    tokens.push(TriviaToken {
+                        start: *prev_start,
+                        end: *prev_end,
+                        kind: TriviaTokenKind::MagicTrailingColon,
+                    });
                 }
             }
         }
@@ -844,7 +856,7 @@ pub fn decorate_trivia(tokens: Vec<TriviaToken>, python_ast: &[Stmt]) -> TriviaI
                     unreachable!("Attach token to the ast: {:?}", token);
                 }
             }
-            TriviaTokenKind::MagicTrailingComma => {
+            TriviaTokenKind::MagicTrailingComma | TriviaTokenKind::MagicTrailingColon => {
                 if let Some(enclosing_node) = enclosing_node {
                     add_comment(
                         Trivia::from_token(&token, Relationship::Trailing),