Handle end-of-line comments on excepthandler and alias (#3196)

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<ASTFormatContext<'_>> 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(())
     }
 }

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<ASTFormatContext<'_>> for Excepthandler {
 
 impl Format<ASTFormatContext<'_>> for FormatExcepthandler<'_> {
     fn fmt(&self, f: &mut Formatter<ASTFormatContext>) -> 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<ASTFormatContext<'_>> 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(())

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.

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
 )