diff --git a/crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/slices.py b/crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/slices.py index 165117cdcb..63885bb872 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/slices.py +++ b/crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/slices.py @@ -29,3 +29,35 @@ 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] + +slice[::, ::] +slice[ + # A + : + # B + : + # C +] +slice[ + # A + 1: + # B + 2: + # C + 3 +] + +slice[ + # A + 1 + + 2 : + # B + 3 : + # C + 4 +] +x[ + 1: # A + 2: # B + 3 # C +] diff --git a/crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/slices.py.expect b/crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/slices.py.expect index 165117cdcb..61218829e4 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/slices.py.expect +++ b/crates/ruff_python_formatter/resources/test/fixtures/black/simple_cases/slices.py.expect @@ -29,3 +29,31 @@ 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] + +slice[::, ::] +slice[ + # A + : + # B + : + # C +] +slice[ + # A + 1: + # B + 2: + # C + 3 +] + +slice[ + # A + 1 + + 2 : + # B + 3 : + # C + 4 +] +x[1:2:3] # A # B # C diff --git a/crates/ruff_python_formatter/src/attachment.rs b/crates/ruff_python_formatter/src/attachment.rs index 8c52cfbf0c..233f1006ea 100644 --- a/crates/ruff_python_formatter/src/attachment.rs +++ b/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) { diff --git a/crates/ruff_python_formatter/src/core/visitor.rs b/crates/ruff_python_formatter/src/core/visitor.rs index cfb07a6210..51fc640a97 100644 --- a/crates/ruff_python_formatter/src/core/visitor.rs +++ b/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); diff --git a/crates/ruff_python_formatter/src/cst.rs b/crates/ruff_python_formatter/src/cst.rs index 3a620bdaec..b193da3d06 100644 --- a/crates/ruff_python_formatter/src/cst.rs +++ b/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>, - upper: Option>, - step: Option>, + lower: SliceIndex, + upper: SliceIndex, + step: Option, }, } @@ -421,6 +423,16 @@ pub enum ExcepthandlerKind { pub type Excepthandler = Located; +#[derive(Clone, Debug, PartialEq)] +pub enum SliceIndexKind { + /// The index slot exists, but is empty. + Empty, + /// The index slot contains an expression. + Index { value: Box }, +} + +pub type SliceIndex = Located; + #[derive(Clone, Debug, PartialEq)] pub struct Arguments { pub posonlyargs: Vec, @@ -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, + } + } } } } diff --git a/crates/ruff_python_formatter/src/format/expr.rs b/crates/ruff_python_formatter/src/format/expr.rs index ee51970e56..3c1dd93fc2 100644 --- a/crates/ruff_python_formatter/src/format/expr.rs +++ b/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>, 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> 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) } diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap index 34ab0f3c2b..adb746d4f8 100644 --- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__expression_py.snap +++ b/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"} diff --git a/crates/ruff_python_formatter/src/trivia.rs b/crates/ruff_python_formatter/src/trivia.rs index 7c5f2f1c46..7eb9c03fd8 100644 --- a/crates/ruff_python_formatter/src/trivia.rs +++ b/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 { 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>) { 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>) { } } 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>) { 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>, pub alias: FxHashMap>, pub excepthandler: FxHashMap>, - pub withitem: FxHashMap>, + pub slice_index: FxHashMap>, } 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, 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),