From 869807fc211479b0059fce2a73f219900ea29075 Mon Sep 17 00:00:00 2001 From: Micha Reiser Date: Tue, 17 Sep 2024 12:30:01 +0200 Subject: [PATCH] Correctly handle other parameter comment placements --- .../test/fixtures/ruff/statement/function.py | 35 +++++++++ .../src/comments/placement.rs | 52 ++++++++++---- .../src/other/parameter.rs | 23 ++++-- .../format@statement__function.py.snap | 71 +++++++++++++++++++ 4 files changed, 164 insertions(+), 17 deletions(-) diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py index 8cebc63d1186d9..3aacd5926d8438 100644 --- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py +++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/function.py @@ -528,3 +528,38 @@ def kwargs_type_annotation(** # comment kwargs: int): pass + +def args_many_comments( + # before + * + # between * and name + args # trailing args + # after name + ): pass + + +def args_many_comments_with_type_annotation( + # before + * + # between * and name + args # trailing args + # before colon + : # after colon + # before type + int # trailing type + # after type + ): pass + + + +def args_with_type_annotations_no_after_colon_comment( + # before + * + # between * and name + args # trailing args + # before colon + : + # before type + int # trailing type + # after type + ): pass diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs index ac86f24c6adcc5..931fb1420b2842 100644 --- a/crates/ruff_python_formatter/src/comments/placement.rs +++ b/crates/ruff_python_formatter/src/comments/placement.rs @@ -2,7 +2,9 @@ use std::cmp::Ordering; use ast::helpers::comment_indentation_after; use ruff_python_ast::whitespace::indentation; -use ruff_python_ast::{self as ast, AnyNodeRef, Comprehension, Expr, ModModule, Parameters}; +use ruff_python_ast::{ + self as ast, AnyNodeRef, Comprehension, Expr, ModModule, Parameter, Parameters, +}; use ruff_python_trivia::{ find_only_token_in_range, indentation_at_offset, BackwardsTokenizer, CommentRanges, SimpleToken, SimpleTokenKind, SimpleTokenizer, @@ -202,17 +204,7 @@ fn handle_enclosed_comment<'a>( } }) } - AnyNodeRef::Parameter(parameter) => { - let is_arg_or_kwarg = locator.slice(parameter.range()).starts_with('*'); - - // A comment between the `*` or `**` and the parameter name. - // Move the comment before the node. - if is_arg_or_kwarg && comment.preceding_node().is_none() { - CommentPlacement::leading(parameter, comment) - } else { - CommentPlacement::Default(comment) - } - } + AnyNodeRef::Parameter(parameter) => handle_parameter_comment(comment, parameter, locator), AnyNodeRef::Arguments(_) | AnyNodeRef::TypeParams(_) | AnyNodeRef::PatternArguments(_) => { handle_bracketed_end_of_line_comment(comment, locator) } @@ -763,6 +755,42 @@ fn handle_parameters_separator_comment<'a>( CommentPlacement::Default(comment) } +/// Associate comments that come before the `:` starting the type annotation or before the +/// parameter's name for unannotated parameters as leading parameter-comments. +/// +/// The parameter's name isn't a node to which comments can be associated. +/// That's why we pull out all comments that come before the expression name or the type annotation +/// and make them leading parameter comments. For example: +/// * `* # comment\nargs` +/// * `arg # comment\n : int` +/// +/// Associate comments with the type annotation when possible. +fn handle_parameter_comment<'a>( + comment: DecoratedComment<'a>, + parameter: &'a Parameter, + locator: &Locator, +) -> CommentPlacement<'a> { + if parameter.annotation.as_deref().is_some() { + let colon = SimpleTokenizer::starts_at(parameter.name.end(), locator.contents()) + .skip_trivia() + .next() + .expect("A annotated parameter should have a colon following its name when it is valid syntax."); + + assert_eq!(colon.kind(), SimpleTokenKind::Colon); + + if comment.start() < colon.start() { + // The comment is before the colon, pull it out and make it a leading comment of the parameter. + CommentPlacement::leading(parameter, comment) + } else { + CommentPlacement::Default(comment) + } + } else if comment.start() < parameter.name.start() { + CommentPlacement::leading(parameter, comment) + } else { + CommentPlacement::Default(comment) + } +} + /// Handles comments between the left side and the operator of a binary expression (trailing comments of the left), /// and trailing end-of-line comments that are on the same line as the operator. /// diff --git a/crates/ruff_python_formatter/src/other/parameter.rs b/crates/ruff_python_formatter/src/other/parameter.rs index fb0e84fbcd46dd..8a634928fcf6da 100644 --- a/crates/ruff_python_formatter/src/other/parameter.rs +++ b/crates/ruff_python_formatter/src/other/parameter.rs @@ -1,7 +1,6 @@ -use ruff_formatter::write; -use ruff_python_ast::Parameter; - +use crate::expression::parentheses::is_expression_parenthesized; use crate::prelude::*; +use ruff_python_ast::Parameter; #[derive(Default)] pub struct FormatParameter; @@ -16,8 +15,22 @@ impl FormatNodeRule for FormatParameter { name.format().fmt(f)?; - if let Some(annotation) = annotation { - write!(f, [token(":"), space(), annotation.format()])?; + if let Some(annotation) = annotation.as_deref() { + token(":").fmt(f)?; + + if f.context().comments().has_leading(annotation) + && !is_expression_parenthesized( + annotation.into(), + f.context().comments().ranges(), + f.context().source(), + ) + { + hard_line_break().fmt(f)?; + } else { + space().fmt(f)?; + } + + annotation.format().fmt(f)?; } Ok(()) diff --git a/crates/ruff_python_formatter/tests/snapshots/format@statement__function.py.snap b/crates/ruff_python_formatter/tests/snapshots/format@statement__function.py.snap index aef8f8d45c5ba0..64eaa01b0b974c 100644 --- a/crates/ruff_python_formatter/tests/snapshots/format@statement__function.py.snap +++ b/crates/ruff_python_formatter/tests/snapshots/format@statement__function.py.snap @@ -534,6 +534,41 @@ def kwargs_type_annotation(** # comment kwargs: int): pass + +def args_many_comments( + # before + * + # between * and name + args # trailing args + # after name + ): pass + + +def args_many_comments_with_type_annotation( + # before + * + # between * and name + args # trailing args + # before colon + : # after colon + # before type + int # trailing type + # after type + ): pass + + + +def args_with_type_annotations_no_after_colon_comment( + # before + * + # between * and name + args # trailing args + # before colon + : + # before type + int # trailing type + # after type + ): pass ``` ## Output @@ -1247,6 +1282,42 @@ def kwargs_type_annotation( **kwargs: int, ): pass + + +def args_many_comments( + # before + # between * and name + *args, # trailing args + # after name +): + pass + + +def args_many_comments_with_type_annotation( + # before + # between * and name + # trailing args + # before colon + *args: + # after colon + # before type + int, # trailing type + # after type +): + pass + + +def args_with_type_annotations_no_after_colon_comment( + # before + # between * and name + # trailing args + # before colon + *args: + # before type + int, # trailing type + # after type +): + pass ```