diff --git a/crates/ruff_python_formatter/src/format/alias.rs b/crates/ruff_python_formatter/src/format/alias.rs index 5ed0b47d08..c0ae052a10 100644 --- a/crates/ruff_python_formatter/src/format/alias.rs +++ b/crates/ruff_python_formatter/src/format/alias.rs @@ -2,9 +2,11 @@ use ruff_formatter::prelude::*; use ruff_formatter::write; use ruff_text_size::TextSize; +use crate::builders::literal; use crate::context::ASTFormatContext; use crate::cst::Alias; use crate::shared_traits::AsFormat; +use crate::trivia::{Relationship, TriviaKind}; pub struct FormatAlias<'a> { item: &'a Alias, @@ -28,6 +30,25 @@ impl Format> for FormatAlias<'_> { write!(f, [dynamic_text(asname, TextSize::default())])?; } + // Format any end-of-line comments. + let mut first = true; + for range in alias.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(()) } } diff --git a/crates/ruff_python_formatter/src/format/excepthandler.rs b/crates/ruff_python_formatter/src/format/excepthandler.rs index 8c59d8ee2b..dfcb688f6d 100644 --- a/crates/ruff_python_formatter/src/format/excepthandler.rs +++ b/crates/ruff_python_formatter/src/format/excepthandler.rs @@ -2,10 +2,12 @@ use ruff_formatter::prelude::*; use ruff_formatter::write; use ruff_text_size::TextSize; +use crate::builders::literal; use crate::context::ASTFormatContext; use crate::cst::{Excepthandler, ExcepthandlerKind}; use crate::format::builders::block; use crate::shared_traits::AsFormat; +use crate::trivia::{Relationship, TriviaKind}; pub struct FormatExcepthandler<'a> { item: &'a Excepthandler, @@ -21,7 +23,8 @@ impl AsFormat> for Excepthandler { impl Format> for FormatExcepthandler<'_> { fn fmt(&self, f: &mut Formatter) -> FormatResult<()> { - let ExcepthandlerKind::ExceptHandler { type_, name, body } = &self.item.node; + let excepthandler = self.item; + let ExcepthandlerKind::ExceptHandler { type_, name, body } = &excepthandler.node; write!(f, [text("except")])?; if let Some(type_) = &type_ { @@ -39,6 +42,26 @@ impl Format> for FormatExcepthandler<'_> { } } write!(f, [text(":")])?; + + // Format any end-of-line comments. + let mut first = true; + for range in excepthandler.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))])?; + } + write!(f, [block_indent(&block(body))])?; Ok(()) diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments2_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments2_py.snap index fb232c5523..a0f5f2619b 100644 --- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments2_py.snap +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments2_py.snap @@ -178,17 +178,6 @@ instruction()#comment with bad spacing ```diff --- Black +++ Ruff -@@ -1,8 +1,8 @@ - from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import ( -- MyLovelyCompanyTeamProjectComponent, # NOT DRY -+ MyLovelyCompanyTeamProjectComponent, - ) - from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import ( -- MyLovelyCompanyTeamProjectComponent as component, # DRY -+ MyLovelyCompanyTeamProjectComponent as component, - ) - - # Please keep __all__ alphabetized within each category. @@ -13,7 +13,7 @@ "Callable", "ClassVar", @@ -333,10 +322,10 @@ instruction()#comment with bad spacing ```py from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import ( - MyLovelyCompanyTeamProjectComponent, + MyLovelyCompanyTeamProjectComponent, # NOT DRY ) from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import ( - MyLovelyCompanyTeamProjectComponent as component, + MyLovelyCompanyTeamProjectComponent as component, # DRY ) # Please keep __all__ alphabetized within each category. diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments4_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments4_py.snap index c5e260faf6..662947d62f 100644 --- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments4_py.snap +++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments4_py.snap @@ -107,17 +107,6 @@ def foo3(list_a, list_b): ```diff --- Black +++ Ruff -@@ -1,8 +1,8 @@ - from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import ( -- MyLovelyCompanyTeamProjectComponent, # NOT DRY -+ MyLovelyCompanyTeamProjectComponent, - ) - from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import ( -- MyLovelyCompanyTeamProjectComponent as component, # DRY -+ MyLovelyCompanyTeamProjectComponent as component, - ) - - @@ -10,39 +10,50 @@ @pytest.mark.parametrize( ("post_data", "message"), @@ -249,10 +238,10 @@ def foo3(list_a, list_b): ```py from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import ( - MyLovelyCompanyTeamProjectComponent, + MyLovelyCompanyTeamProjectComponent, # NOT DRY ) from com.my_lovely_company.my_lovely_team.my_lovely_project.my_lovely_component import ( - MyLovelyCompanyTeamProjectComponent as component, + MyLovelyCompanyTeamProjectComponent as component, # DRY )