Introduce a new CST element for slice segments (#3195)

crates/ruff_python_formatter/src/attachment.rs

@@ -1,6 +1,6 @@
 use crate::core::visitor;
 use crate::core::visitor::Visitor;
-use crate::cst::{Expr, Stmt};
+use crate::cst::{Alias, Excepthandler, Expr, SliceIndex, Stmt};
 use crate::trivia::{decorate_trivia, TriviaIndex, TriviaToken};
 
 struct AttachmentVisitor {
@@ -23,6 +23,30 @@ impl<'a> Visitor<'a> for AttachmentVisitor {
         }
         visitor::walk_expr(self, expr);
     }
+
+    fn visit_alias(&mut self, alias: &'a mut Alias) {
+        let trivia = self.index.alias.remove(&alias.id());
+        if let Some(comments) = trivia {
+            alias.trivia.extend(comments);
+        }
+        visitor::walk_alias(self, alias);
+    }
+
+    fn visit_excepthandler(&mut self, excepthandler: &'a mut Excepthandler) {
+        let trivia = self.index.excepthandler.remove(&excepthandler.id());
+        if let Some(comments) = trivia {
+            excepthandler.trivia.extend(comments);
+        }
+        visitor::walk_excepthandler(self, excepthandler);
+    }
+
+    fn visit_slice_index(&mut self, slice_index: &'a mut SliceIndex) {
+        let trivia = self.index.slice_index.remove(&slice_index.id());
+        if let Some(comments) = trivia {
+            slice_index.trivia.extend(comments);
+        }
+        visitor::walk_slice_index(self, slice_index);
+    }
 }
 
 pub fn attach(python_cst: &mut [Stmt], trivia: Vec<TriviaToken>) {

crates/ruff_python_formatter/src/core/visitor.rs

@@ -2,8 +2,8 @@ use rustpython_parser::ast::Constant;
 
 use crate::cst::{
     Alias, Arg, Arguments, Boolop, Cmpop, Comprehension, Excepthandler, ExcepthandlerKind, Expr,
-    ExprContext, ExprKind, Keyword, MatchCase, Operator, Pattern, PatternKind, Stmt, StmtKind,
-    Unaryop, Withitem,
+    ExprContext, ExprKind, Keyword, MatchCase, Operator, Pattern, PatternKind, SliceIndex,
+    SliceIndexKind, Stmt, StmtKind, Unaryop, Withitem,
 };
 
 pub trait Visitor<'a> {
@@ -40,6 +40,9 @@ pub trait Visitor<'a> {
     fn visit_excepthandler(&mut self, excepthandler: &'a mut Excepthandler) {
         walk_excepthandler(self, excepthandler);
     }
+    fn visit_slice_index(&mut self, slice_index: &'a mut SliceIndex) {
+        walk_slice_index(self, slice_index);
+    }
     fn visit_format_spec(&mut self, format_spec: &'a mut Expr) {
         walk_expr(self, format_spec);
     }
@@ -420,14 +423,10 @@ pub fn walk_expr<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, expr: &'a mut Exp
             visitor.visit_expr_context(ctx);
         }
         ExprKind::Slice { lower, upper, step } => {
-            if let Some(expr) = lower {
-                visitor.visit_expr(expr);
-            }
-            if let Some(expr) = upper {
-                visitor.visit_expr(expr);
-            }
+            visitor.visit_slice_index(lower);
+            visitor.visit_slice_index(upper);
             if let Some(expr) = step {
-                visitor.visit_expr(expr);
+                visitor.visit_slice_index(expr);
             }
         }
     }
@@ -466,6 +465,18 @@ pub fn walk_excepthandler<'a, V: Visitor<'a> + ?Sized>(
     }
 }
 
+pub fn walk_slice_index<'a, V: Visitor<'a> + ?Sized>(
+    visitor: &mut V,
+    slice_index: &'a mut SliceIndex,
+) {
+    match &mut slice_index.node {
+        SliceIndexKind::Empty => {}
+        SliceIndexKind::Index { value } => {
+            visitor.visit_expr(value);
+        }
+    }
+}
+
 pub fn walk_arguments<'a, V: Visitor<'a> + ?Sized>(visitor: &mut V, arguments: &'a mut Arguments) {
     for arg in &mut arguments.posonlyargs {
         visitor.visit_arg(arg);

crates/ruff_python_formatter/src/cst.rs

@@ -1,8 +1,10 @@
 #![allow(clippy::derive_partial_eq_without_eq)]
 
 use rustpython_parser::ast::{Constant, Location};
+use rustpython_parser::Mode;
 
 use crate::core::locator::Locator;
+use crate::core::types::Range;
 use crate::trivia::{Parenthesize, Trivia};
 
 type Ident = String;
@@ -394,9 +396,9 @@ pub enum ExprKind {
         ctx: ExprContext,
     },
     Slice {
-        lower: Option<Box<Expr>>,
-        upper: Option<Box<Expr>>,
-        step: Option<Box<Expr>>,
+        lower: SliceIndex,
+        upper: SliceIndex,
+        step: Option<SliceIndex>,
     },
 }
 
@@ -421,6 +423,16 @@ pub enum ExcepthandlerKind {
 
 pub type Excepthandler = Located<ExcepthandlerKind>;
 
+#[derive(Clone, Debug, PartialEq)]
+pub enum SliceIndexKind {
+    /// The index slot exists, but is empty.
+    Empty,
+    /// The index slot contains an expression.
+    Index { value: Box<Expr> },
+}
+
+pub type SliceIndex = Located<SliceIndexKind>;
+
 #[derive(Clone, Debug, PartialEq)]
 pub struct Arguments {
     pub posonlyargs: Vec<Arg>,
@@ -1404,17 +1416,79 @@ impl From<(rustpython_parser::ast::Expr, &Locator<'_>)> for Expr {
                 trivia: vec![],
                 parentheses: Parenthesize::Never,
             },
-            rustpython_parser::ast::ExprKind::Slice { lower, upper, step } => Expr {
-                location: expr.location,
-                end_location: expr.end_location,
-                node: ExprKind::Slice {
-                    lower: lower.map(|node| Box::new((*node, locator).into())),
-                    upper: upper.map(|node| Box::new((*node, locator).into())),
-                    step: step.map(|node| Box::new((*node, locator).into())),
-                },
-                trivia: vec![],
-                parentheses: Parenthesize::Never,
-            },
+            rustpython_parser::ast::ExprKind::Slice { lower, upper, step } => {
+                // Locate the colon tokens, which indicate the number of index segments.
+                let (source, start, end) =
+                    locator.slice(Range::new(expr.location, expr.end_location.unwrap()));
+                let tokens = rustpython_parser::lexer::lex_located(
+                    &source[start..end],
+                    Mode::Module,
+                    expr.location,
+                );
+
+                // Find the first and (if it exists) second colon in the slice, avoiding any
+                // semicolons within nested slices, and any lambda expressions.
+                let mut first_colon = None;
+                let mut second_colon = None;
+                let mut lambda = 0;
+                let mut nesting = 0;
+                for (start, tok, ..) in tokens.flatten() {
+                    match tok {
+                        rustpython_parser::Tok::Lambda if nesting == 0 => lambda += 1,
+                        rustpython_parser::Tok::Colon if nesting == 0 => {
+                            if lambda > 0 {
+                                lambda -= 1;
+                            } else {
+                                if first_colon.is_none() {
+                                    first_colon = Some(start);
+                                } else {
+                                    second_colon = Some(start);
+                                    break;
+                                }
+                            }
+                        }
+                        rustpython_parser::Tok::Lpar
+                        | rustpython_parser::Tok::Lsqb
+                        | rustpython_parser::Tok::Lbrace => nesting += 1,
+                        rustpython_parser::Tok::Rpar
+                        | rustpython_parser::Tok::Rsqb
+                        | rustpython_parser::Tok::Rbrace => nesting -= 1,
+                        _ => {}
+                    }
+                }
+
+                let lower = SliceIndex::new(
+                    expr.location,
+                    first_colon.unwrap(),
+                    lower.map_or(SliceIndexKind::Empty, |node| SliceIndexKind::Index {
+                        value: Box::new((*node, locator).into()),
+                    }),
+                );
+                let upper = SliceIndex::new(
+                    first_colon.unwrap(),
+                    second_colon.unwrap_or(expr.end_location.unwrap()),
+                    upper.map_or(SliceIndexKind::Empty, |node| SliceIndexKind::Index {
+                        value: Box::new((*node, locator).into()),
+                    }),
+                );
+                let step = second_colon.map(|second_colon| {
+                    SliceIndex::new(
+                        second_colon,
+                        expr.end_location.unwrap(),
+                        step.map_or(SliceIndexKind::Empty, |node| SliceIndexKind::Index {
+                            value: Box::new((*node, locator).into()),
+                        }),
+                    )
+                });
+
+                Expr {
+                    location: expr.location,
+                    end_location: expr.end_location,
+                    node: ExprKind::Slice { lower, upper, step },
+                    trivia: vec![],
+                    parentheses: Parenthesize::Never,
+                }
+            }
         }
     }
 }

crates/ruff_python_formatter/src/format/expr.rs

@@ -10,7 +10,8 @@ use crate::builders::literal;
 use crate::context::ASTFormatContext;
 use crate::core::types::Range;
 use crate::cst::{
-    Arguments, Boolop, Cmpop, Comprehension, Expr, ExprKind, Keyword, Operator, Unaryop,
+    Arguments, Boolop, Cmpop, Comprehension, Expr, ExprKind, Keyword, Operator, SliceIndex,
+    SliceIndexKind, Unaryop,
 };
 use crate::format::helpers::{is_self_closing, is_simple_power, is_simple_slice};
 use crate::format::numbers::{complex_literal, float_literal, int_literal};
@@ -86,15 +87,30 @@ fn format_subscript(
     value: &Expr,
     slice: &Expr,
 ) -> FormatResult<()> {
+    write!(f, [value.format()])?;
+    write!(f, [text("[")])?;
     write!(
         f,
-        [
-            value.format(),
-            text("["),
-            group(&format_args![soft_block_indent(&slice.format())]),
-            text("]")
-        ]
+        [group(&format_args![soft_block_indent(&format_with(|f| {
+            write!(f, [slice.format()])?;
+
+            // Apply any dangling comments.
+            for trivia in &expr.trivia {
+                if matches!(trivia.relationship, Relationship::Dangling) {
+                    if let TriviaKind::OwnLineComment(range) = trivia.kind {
+                        write!(f, [expand_parent()])?;
+                        write!(f, [hard_line_break()])?;
+                        write!(f, [literal(range)])?;
+                    }
+                }
+            }
+
+            Ok(())
+        }))])]
     )?;
+
+    write!(f, [text("]")])?;
+
     Ok(())
 }
 
@@ -187,51 +203,182 @@ fn format_tuple(
 fn format_slice(
     f: &mut Formatter<ASTFormatContext<'_>>,
     expr: &Expr,
-    lower: Option<&Expr>,
-    upper: Option<&Expr>,
-    step: Option<&Expr>,
+    lower: &SliceIndex,
+    upper: &SliceIndex,
+    step: Option<&SliceIndex>,
 ) -> FormatResult<()> {
-    // https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#slices
-    let is_simple = lower.map_or(true, is_simple_slice)
-        && upper.map_or(true, is_simple_slice)
-        && step.map_or(true, is_simple_slice);
-
-    if let Some(lower) = lower {
-        write!(f, [lower.format()])?;
-        if !is_simple {
-            write!(f, [space()])?;
-        }
-    }
-    write!(f, [text(":")])?;
-    if let Some(upper) = upper {
-        if !is_simple {
-            write!(f, [space()])?;
-        }
-        write!(f, [upper.format()])?;
-        if !is_simple && step.is_some() {
-            write!(f, [space()])?;
-        }
-    }
-    if let Some(step) = step {
-        if !is_simple && upper.is_some() {
-            write!(f, [space()])?;
-        }
-        write!(f, [text(":")])?;
-        if !is_simple {
-            write!(f, [space()])?;
-        }
-        write!(f, [step.format()])?;
+    // // https://black.readthedocs.io/en/stable/the_black_code_style/current_style.html#slices
+    let lower_is_simple = if let SliceIndexKind::Index { value } = &lower.node {
+        is_simple_slice(value)
     } 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()])?;
+        true
+    };
+    let upper_is_simple = if let SliceIndexKind::Index { value } = &upper.node {
+        is_simple_slice(value)
+    } else {
+        true
+    };
+    let step_is_simple = step.map_or(true, |step| {
+        if let SliceIndexKind::Index { value } = &step.node {
+            is_simple_slice(value)
+        } else {
+            true
+        }
+    });
+    let is_simple = lower_is_simple && upper_is_simple && step_is_simple;
+
+    write!(
+        f,
+        [group(&format_with(|f| {
+            if let SliceIndexKind::Index { value } = &lower.node {
+                write!(f, [value.format()])?;
+            }
+
+            // Apply any dangling comments.
+            for trivia in &lower.trivia {
+                if matches!(trivia.relationship, Relationship::Dangling) {
+                    if let TriviaKind::OwnLineComment(range) = trivia.kind {
+                        write!(f, [expand_parent()])?;
+                        write!(f, [hard_line_break()])?;
+                        write!(f, [literal(range)])?;
+                        write!(f, [hard_line_break()])?;
+                    }
+                }
+            }
+
+            if matches!(lower.node, SliceIndexKind::Index { .. }) {
+                if !is_simple {
+                    write!(f, [space()])?;
+                }
             }
             write!(f, [text(":")])?;
+
+            // Format any end-of-line comments.
+            let mut first = true;
+            for range in lower.trivia.iter().filter_map(|trivia| {
+                if matches!(trivia.relationship, Relationship::Trailing) {
+                    if let TriviaKind::EndOfLineComment(range) = trivia.kind {
+                        Some(range)
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            }) {
+                if std::mem::take(&mut first) {
+                    write!(f, [line_suffix(&text("  "))])?;
+                }
+                write!(f, [line_suffix(&literal(range))])?;
+            }
+
+            if let SliceIndexKind::Index { value } = &upper.node {
+                if !is_simple {
+                    write!(f, [space()])?;
+                }
+                write!(f, [if_group_breaks(&soft_line_break())])?;
+                write!(f, [value.format()])?;
+            }
+
+            // Apply any dangling comments.
+            for trivia in &upper.trivia {
+                if matches!(trivia.relationship, Relationship::Dangling) {
+                    if let TriviaKind::OwnLineComment(range) = trivia.kind {
+                        write!(f, [expand_parent()])?;
+                        write!(f, [hard_line_break()])?;
+                        write!(f, [literal(range)])?;
+                        write!(f, [hard_line_break()])?;
+                    }
+                }
+            }
+
+            // Format any end-of-line comments.
+            let mut first = true;
+            for range in upper.trivia.iter().filter_map(|trivia| {
+                if matches!(trivia.relationship, Relationship::Trailing) {
+                    if let TriviaKind::EndOfLineComment(range) = trivia.kind {
+                        Some(range)
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            }) {
+                if std::mem::take(&mut first) {
+                    write!(f, [line_suffix(&text("  "))])?;
+                }
+                write!(f, [line_suffix(&literal(range))])?;
+            }
+
+            if let Some(step) = step {
+                if matches!(upper.node, SliceIndexKind::Index { .. }) {
+                    if !is_simple {
+                        write!(f, [space()])?;
+                    }
+                }
+                write!(f, [text(":")])?;
+
+                if let SliceIndexKind::Index { value } = &step.node {
+                    if !is_simple {
+                        write!(f, [space()])?;
+                    }
+                    write!(f, [if_group_breaks(&soft_line_break())])?;
+                    write!(f, [value.format()])?;
+                }
+
+                // Apply any dangling comments.
+                for trivia in &step.trivia {
+                    if matches!(trivia.relationship, Relationship::Dangling) {
+                        if let TriviaKind::OwnLineComment(range) = trivia.kind {
+                            write!(f, [expand_parent()])?;
+                            write!(f, [hard_line_break()])?;
+                            write!(f, [literal(range)])?;
+                            write!(f, [hard_line_break()])?;
+                        }
+                    }
+                }
+
+                // Format any end-of-line comments.
+                let mut first = true;
+                for range in step.trivia.iter().filter_map(|trivia| {
+                    if matches!(trivia.relationship, Relationship::Trailing) {
+                        if let TriviaKind::EndOfLineComment(range) = trivia.kind {
+                            Some(range)
+                        } else {
+                            None
+                        }
+                    } else {
+                        None
+                    }
+                }) {
+                    if std::mem::take(&mut first) {
+                        write!(f, [line_suffix(&text("  "))])?;
+                    }
+                    write!(f, [line_suffix(&literal(range))])?;
+                }
+            }
+            Ok(())
+        }))]
+    )?;
+
+    // Format any end-of-line comments.
+    let mut first = true;
+    for range in expr.trivia.iter().filter_map(|trivia| {
+        if matches!(trivia.relationship, Relationship::Trailing) {
+            if let TriviaKind::EndOfLineComment(range) = trivia.kind {
+                Some(range)
+            } else {
+                None
+            }
+        } else {
+            None
         }
+    }) {
+        if std::mem::take(&mut first) {
+            write!(f, [line_suffix(&text("  "))])?;
+        }
+        write!(f, [line_suffix(&literal(range))])?;
     }
 
     Ok(())
@@ -973,13 +1120,9 @@ impl Format<ASTFormatContext<'_>> for FormatExpr<'_> {
             ExprKind::Name { id, .. } => format_name(f, self.item, id),
             ExprKind::List { elts, .. } => format_list(f, self.item, elts),
             ExprKind::Tuple { elts, .. } => format_tuple(f, self.item, elts),
-            ExprKind::Slice { lower, upper, step } => format_slice(
-                f,
-                self.item,
-                lower.as_deref(),
-                upper.as_deref(),
-                step.as_deref(),
-            ),
+            ExprKind::Slice { lower, upper, step } => {
+                format_slice(f, self.item, lower, upper, step.as_ref())
+            }
             _ => {
                 unimplemented!("Implement ExprKind: {:?}", self.item.node)
             }

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap

@@ -292,7 +292,7 @@ last_call()
  }
  Python3 > Python2 > COBOL
  Life is Life
-@@ -138,33 +141,33 @@
+@@ -138,15 +141,15 @@
  very_long_variable_name_filters: t.List[
      t.Tuple[str, t.Union[str, t.List[t.Optional[str]]]],
  ]
@@ -315,37 +315,6 @@ last_call()
  slice[0]
  slice[0:1]
  slice[0:1:2]
--slice[:]
-+slice[::]
- slice[:-1]
--slice[1:]
-+slice[1::]
- slice[::-1]
- slice[d :: d + 1]
- slice[:c, c - 1]
- numpy[:, 0:1]
- numpy[:, :-1]
--numpy[0, :]
-+numpy[0, ::]
- numpy[:, i]
- numpy[0, :2]
- numpy[:N, 0]
- numpy[:2, :4]
- numpy[2:4, 1:5]
--numpy[4:, 2:]
-+numpy[4:, 2::]
- numpy[:, (0, 1, 2, 5)]
- numpy[0, [0]]
- numpy[:, [i]]
-@@ -172,7 +175,7 @@
- numpy[-(c + 1) :, d]
- numpy[:, l[-2]]
- numpy[:, ::-1]
--numpy[np.newaxis, :]
-+numpy[np.newaxis, ::]
- (str or None) if (sys.version_info[0] > (3,)) else (str or bytes or None)
- {"2.7": dead, "3.7": long_live or die_hard}
- {"2.7", "3.6", "3.7", "3.8", "3.9", "4.0" if gilectomy else "3.10"}
 @@ -201,30 +204,26 @@
  e = (1,).count(1)
  f = 1, *range(10)
@@ -600,21 +569,21 @@ xxxx_xxx_xxxx_xxxxx_xxxx_xxx: Callable[..., List[SomeClass]] = (
 slice[0]
 slice[0:1]
 slice[0:1:2]
-slice[::]
+slice[:]
 slice[:-1]
-slice[1::]
+slice[1:]
 slice[::-1]
 slice[d :: d + 1]
 slice[:c, c - 1]
 numpy[:, 0:1]
 numpy[:, :-1]
-numpy[0, ::]
+numpy[0, :]
 numpy[:, i]
 numpy[0, :2]
 numpy[:N, 0]
 numpy[:2, :4]
 numpy[2:4, 1:5]
-numpy[4:, 2::]
+numpy[4:, 2:]
 numpy[:, (0, 1, 2, 5)]
 numpy[0, [0]]
 numpy[:, [i]]
@@ -622,7 +591,7 @@ numpy[1 : c + 1, c]
 numpy[-(c + 1) :, d]
 numpy[:, l[-2]]
 numpy[:, ::-1]
-numpy[np.newaxis, ::]
+numpy[np.newaxis, :]
 (str or None) if (sys.version_info[0] > (3,)) else (str or bytes or None)
 {"2.7": dead, "3.7": long_live or die_hard}
 {"2.7", "3.6", "3.7", "3.8", "3.9", "4.0" if gilectomy else "3.10"}

crates/ruff_python_formatter/src/trivia.rs

@@ -4,7 +4,10 @@ use rustpython_parser::lexer::LexResult;
 use rustpython_parser::Tok;
 
 use crate::core::types::Range;
-use crate::cst::{Alias, Excepthandler, ExcepthandlerKind, Expr, ExprKind, Stmt, StmtKind};
+use crate::cst::{
+    Alias, Excepthandler, ExcepthandlerKind, Expr, ExprKind, SliceIndex, SliceIndexKind, Stmt,
+    StmtKind,
+};
 
 #[derive(Clone, Debug)]
 pub enum Node<'a> {
@@ -13,6 +16,7 @@ pub enum Node<'a> {
     Expr(&'a Expr),
     Alias(&'a Alias),
     Excepthandler(&'a Excepthandler),
+    SliceIndex(&'a SliceIndex),
 }
 
 impl Node<'_> {
@@ -23,6 +27,7 @@ impl Node<'_> {
             Node::Expr(node) => node.id(),
             Node::Alias(node) => node.id(),
             Node::Excepthandler(node) => node.id(),
+            Node::SliceIndex(node) => node.id(),
         }
     }
 }
@@ -32,7 +37,6 @@ pub enum TriviaTokenKind {
     OwnLineComment,
     EndOfLineComment,
     MagicTrailingComma,
-    MagicTrailingColon,
     EmptyLine,
     Parentheses,
 }
@@ -69,7 +73,6 @@ pub enum TriviaKind {
     /// ```
     EndOfLineComment(Range),
     MagicTrailingComma,
-    MagicTrailingColon,
     EmptyLine,
     Parentheses,
 }
@@ -104,10 +107,6 @@ impl Trivia {
                 kind: TriviaKind::MagicTrailingComma,
                 relationship,
             },
-            TriviaTokenKind::MagicTrailingColon => Self {
-                kind: TriviaKind::MagicTrailingColon,
-                relationship,
-            },
             TriviaTokenKind::EmptyLine => Self {
                 kind: TriviaKind::EmptyLine,
                 relationship,
@@ -171,12 +170,6 @@ 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,
-                    });
                 }
             }
         }
@@ -450,12 +443,8 @@ fn sorted_child_nodes_inner<'a>(node: &Node<'a>, result: &mut Vec<Node<'a>>) {
                     result.push(Node::Alias(name));
                 }
             }
-            StmtKind::Global { .. } => {
-                // TODO(charlie): Ident, not sure how to handle?
-            }
-            StmtKind::Nonlocal { .. } => {
-                // TODO(charlie): Ident, not sure how to handle?
-            }
+            StmtKind::Global { .. } => {}
+            StmtKind::Nonlocal { .. } => {}
         },
         // TODO(charlie): Actual logic, this doesn't do anything.
         Node::Expr(expr) => match &expr.node {
@@ -617,14 +606,10 @@ fn sorted_child_nodes_inner<'a>(node: &Node<'a>, result: &mut Vec<Node<'a>>) {
                 }
             }
             ExprKind::Slice { lower, upper, step } => {
-                if let Some(lower) = lower {
-                    result.push(Node::Expr(lower));
-                }
-                if let Some(upper) = upper {
-                    result.push(Node::Expr(upper));
-                }
+                result.push(Node::SliceIndex(lower));
+                result.push(Node::SliceIndex(upper));
                 if let Some(step) = step {
-                    result.push(Node::Expr(step));
+                    result.push(Node::SliceIndex(step));
                 }
             }
         },
@@ -639,6 +624,11 @@ fn sorted_child_nodes_inner<'a>(node: &Node<'a>, result: &mut Vec<Node<'a>>) {
                 result.push(Node::Stmt(stmt));
             }
         }
+        Node::SliceIndex(slice_index) => {
+            if let SliceIndexKind::Index { value } = &slice_index.node {
+                result.push(Node::Expr(value));
+            }
+        }
     }
 }
 
@@ -679,6 +669,7 @@ pub fn decorate_token<'a>(
             Node::Expr(node) => node.location,
             Node::Alias(node) => node.location,
             Node::Excepthandler(node) => node.location,
+            Node::SliceIndex(node) => node.location,
             Node::Mod(..) => unreachable!("Node::Mod cannot be a child node"),
         };
         let end = match &child {
@@ -686,6 +677,7 @@ pub fn decorate_token<'a>(
             Node::Expr(node) => node.end_location.unwrap(),
             Node::Alias(node) => node.end_location.unwrap(),
             Node::Excepthandler(node) => node.end_location.unwrap(),
+            Node::SliceIndex(node) => node.end_location.unwrap(),
             Node::Mod(..) => unreachable!("Node::Mod cannot be a child node"),
         };
 
@@ -697,6 +689,7 @@ pub fn decorate_token<'a>(
                 Node::Expr(node) => node.location,
                 Node::Alias(node) => node.location,
                 Node::Excepthandler(node) => node.location,
+                Node::SliceIndex(node) => node.location,
                 Node::Mod(..) => unreachable!("Node::Mod cannot be a child node"),
             };
             let existing_end = match &existing {
@@ -704,6 +697,7 @@ pub fn decorate_token<'a>(
                 Node::Expr(node) => node.end_location.unwrap(),
                 Node::Alias(node) => node.end_location.unwrap(),
                 Node::Excepthandler(node) => node.end_location.unwrap(),
+                Node::SliceIndex(node) => node.end_location.unwrap(),
                 Node::Mod(..) => unreachable!("Node::Mod cannot be a child node"),
             };
             if start == existing_start && end == existing_end {
@@ -763,7 +757,7 @@ pub struct TriviaIndex {
     pub expr: FxHashMap<usize, Vec<Trivia>>,
     pub alias: FxHashMap<usize, Vec<Trivia>>,
     pub excepthandler: FxHashMap<usize, Vec<Trivia>>,
-    pub withitem: FxHashMap<usize, Vec<Trivia>>,
+    pub slice_index: FxHashMap<usize, Vec<Trivia>>,
 }
 
 fn add_comment(comment: Trivia, node: &Node, trivia: &mut TriviaIndex) {
@@ -797,6 +791,13 @@ fn add_comment(comment: Trivia, node: &Node, trivia: &mut TriviaIndex) {
                 .or_insert_with(Vec::new)
                 .push(comment);
         }
+        Node::SliceIndex(node) => {
+            trivia
+                .slice_index
+                .entry(node.id())
+                .or_insert_with(Vec::new)
+                .push(comment);
+        }
     }
 }
 
@@ -871,7 +872,7 @@ pub fn decorate_trivia(tokens: Vec<TriviaToken>, python_ast: &[Stmt]) -> TriviaI
                     unreachable!("Attach token to the ast: {:?}", token);
                 }
             }
-            TriviaTokenKind::MagicTrailingComma | TriviaTokenKind::MagicTrailingColon => {
+            TriviaTokenKind::MagicTrailingComma => {
                 if let Some(enclosing_node) = enclosing_node {
                     add_comment(
                         Trivia::from_token(&token, Relationship::Trailing),