Correctly handle newlines after/before comments

Author: MichaReiser

crates/ruff_python_formatter/resources/test/fixtures/ruff/trivia.py

@@ -0,0 +1,32 @@
+
+# Removes the line above
+
+a = 10  # Keeps the line above
+
+# Separated by one line from `a` and `b`
+
+b = 20
+# Adds two lines after `b`
+class Test:
+    def a(self):
+        pass
+        # trailing comment
+
+# two lines before, one line after
+
+c = 30
+
+while a == 10:
+    ...
+
+    # trailing comment with one line before
+
+# one line before this leading comment
+
+d = 40
+
+while b == 20:
+    ...
+    # no empty line before
+
+e = 50  # one empty line before

crates/ruff_python_formatter/src/builders.rs

@@ -1,7 +1,8 @@
 use crate::context::NodeLevel;
 use crate::prelude::*;
-use crate::trivia::lines_before;
+use crate::trivia::{lines_after, skip_trailing_trivia};
 use ruff_formatter::write;
+use ruff_text_size::TextSize;
 use rustpython_parser::ast::Ranged;
 
 /// Provides Python specific extensions to [`Formatter`].
@@ -26,7 +27,7 @@ impl<'buf, 'ast> PyFormatterExtensions<'ast, 'buf> for PyFormatter<'ast, 'buf> {
 pub(crate) struct JoinNodesBuilder<'fmt, 'ast, 'buf> {
     fmt: &'fmt mut PyFormatter<'ast, 'buf>,
     result: FormatResult<()>,
-    has_elements: bool,
+    last_end: Option<TextSize>,
     node_level: NodeLevel,
 }
 
@@ -35,7 +36,7 @@ impl<'fmt, 'ast, 'buf> JoinNodesBuilder<'fmt, 'ast, 'buf> {
         Self {
             fmt,
             result: Ok(()),
-            has_elements: false,
+            last_end: None,
             node_level: level,
         }
     }
@@ -47,22 +48,43 @@ impl<'fmt, 'ast, 'buf> JoinNodesBuilder<'fmt, 'ast, 'buf> {
         T: Ranged,
     {
         let node_level = self.node_level;
-        let separator = format_with(|f: &mut PyFormatter| match node_level {
-            NodeLevel::TopLevel => match lines_before(node.start(), f.context().contents()) {
-                0 | 1 => hard_line_break().fmt(f),
-                2 => empty_line().fmt(f),
-                _ => write!(f, [empty_line(), empty_line()]),
-            },
-            NodeLevel::CompoundStatement => {
-                match lines_before(node.start(), f.context().contents()) {
-                    0 | 1 => hard_line_break().fmt(f),
-                    _ => empty_line().fmt(f),
-                }
+
+        self.result = self.result.and_then(|_| {
+            if let Some(last_end) = self.last_end.replace(node.end()) {
+                let source = self.fmt.context().contents();
+                let count_lines = |offset| {
+                    // It's necessary to skip any trailing line comment because RustPython doesn't include trailing comments
+                    // in the node's range
+                    // ```python
+                    // a # The range of `a` ends right before this comment
+                    //
+                    // b
+                    // ```
+                    //
+                    // Simply using `lines_after` doesn't work if a statement has a trailing comment because
+                    // it then counts the lines between the statement and the trailing comment, which is
+                    // always 0. This is why it skips any trailing trivia (trivia that's on the same line)
+                    // and counts the lines after.
+                    let after_trailing_trivia = skip_trailing_trivia(offset, source);
+                    lines_after(after_trailing_trivia, source)
+                };
+
+                match node_level {
+                    NodeLevel::TopLevel => match dbg!(count_lines(last_end)) {
+                        0 | 1 => hard_line_break().fmt(self.fmt),
+                        2 => empty_line().fmt(self.fmt),
+                        _ => write!(self.fmt, [empty_line(), empty_line()]),
+                    },
+                    NodeLevel::CompoundStatement => match count_lines(last_end) {
+                        0 | 1 => hard_line_break().fmt(self.fmt),
+                        _ => empty_line().fmt(self.fmt),
+                    },
+                    NodeLevel::Expression => hard_line_break().fmt(self.fmt),
+                }?;
             }
-            NodeLevel::Expression => hard_line_break().fmt(f),
-        });
 
-        self.entry_with_separator(&separator, content);
+            content.fmt(self.fmt)
+        });
     }
 
     /// Writes a sequence of node with their content tuples, inserting the appropriate number of line breaks between any two of them
@@ -98,17 +120,20 @@ impl<'fmt, 'ast, 'buf> JoinNodesBuilder<'fmt, 'ast, 'buf> {
     }
 
     /// Writes a single entry using the specified separator to separate the entry from a previous entry.
-    pub(crate) fn entry_with_separator(
+    pub(crate) fn entry_with_separator<T>(
         &mut self,
         separator: &dyn Format<PyFormatContext<'ast>>,
         content: &dyn Format<PyFormatContext<'ast>>,
-    ) {
+        node: &T,
+    ) where
+        T: Ranged,
+    {
         self.result = self.result.and_then(|_| {
-            if self.has_elements {
+            if self.last_end.is_some() {
                 separator.fmt(self.fmt)?;
             }
 
-            self.has_elements = true;
+            self.last_end = Some(node.end());
 
             content.fmt(self.fmt)
         });

crates/ruff_python_formatter/src/lib.rs

@@ -355,11 +355,12 @@ Formatted twice:
     #[ignore]
     #[test]
     fn quick_test() {
-        let src = r#"
-while True:
-    if something.changed:
-        do.stuff()  # trailing comment
-other
+        let src = r#"AAAAAAAAAAAAA = AAAAAAAAAAAAA  # type: ignore
+
+call_to_some_function_asdf(
+    foo,
+    [AAAAAAAAAAAAAAAAAAAAAAA, AAAAAAAAAAAAAAAAAAAAAAA, AAAAAAAAAAAAAAAAAAAAAAA, BBBBBBBBBBBB],  # type: ignore
+)
 "#;
         // Tokenize once
         let mut tokens = Vec::new();

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__collections_py.snap

@@ -84,24 +84,7 @@ if True:
 ```diff
 --- Black
 +++ Ruff
-@@ -1,75 +1,49 @@
- import core, time, a
- 
- from . import A, B, C
--
- # keeps existing trailing comma
- from foo import (
-     bar,
- )
--
- # also keeps existing structure
- from foo import (
-     baz,
-     qux,
- )
--
- # `as` works as well
- from foo import (
+@@ -18,44 +18,26 @@
      xyzzy as magic,
  )
  
@@ -154,12 +137,12 @@ if True:
 -    "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa wraps %s"
 -    % bar
 -)
--
 +y = {"oneple": (1,),}
 +assert False, ("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa wraps %s" % bar)
+ 
  # looping over a 1-tuple should also not get wrapped
  for x in (1,):
-     pass
+@@ -63,13 +45,9 @@
  for (x,) in (1,), (2,), (3,):
      pass
  
@@ -175,7 +158,7 @@ if True:
  print("foo %r", (foo.bar,))
  
  if True:
-@@ -79,21 +53,15 @@
+@@ -79,21 +57,15 @@
      )
  
  if True:
@@ -210,15 +193,18 @@ if True:
 import core, time, a
 
 from . import A, B, C
+
 # keeps existing trailing comma
 from foo import (
     bar,
 )
+
 # also keeps existing structure
 from foo import (
     baz,
     qux,
 )
+
 # `as` works as well
 from foo import (
     xyzzy as magic,
@@ -244,6 +230,7 @@ nested_long_lines = ["aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa", "bbbbbbbbbbbbbb
 x = {"oneple": (1,)}
 y = {"oneple": (1,),}
 assert False, ("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa wraps %s" % bar)
+
 # looping over a 1-tuple should also not get wrapped
 for x in (1,):
     pass

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtonoff3_py.snap

@@ -30,11 +30,8 @@ x = [
 ```diff
 --- Black
 +++ Ruff
-@@ -10,6 +10,9 @@
-     1, 2,
-     3, 4,
+@@ -12,4 +12,6 @@
  ]
-+
  # fmt: on
  
 -x = [1, 2, 3, 4]
@@ -58,7 +55,6 @@ x = [
     1, 2,
     3, 4,
 ]
-
 # fmt: on
 
 x = [

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__fmtonoff5_py.snap

@@ -97,16 +97,7 @@ elif   unformatted:
 ```diff
 --- Black
 +++ Ruff
-@@ -9,8 +9,6 @@
-             ]  # Includes an formatted indentation.
-     },
- )
--
--
- # Regression test for https://github.com/psf/black/issues/2015.
- run(
-     # fmt: off
-@@ -44,7 +42,7 @@
+@@ -44,7 +44,7 @@
          print ( "This won't be formatted" )
      print ( "This won't be formatted either" )
  else:
@@ -115,7 +106,7 @@ elif   unformatted:
  
  
  # Regression test for https://github.com/psf/black/issues/3184.
-@@ -61,7 +59,7 @@
+@@ -61,7 +61,7 @@
              elif param[0:4] in ("ZZZZ",):
                  print ( "This won't be formatted either" )
  
@@ -124,7 +115,7 @@ elif   unformatted:
  
  
  # Regression test for https://github.com/psf/black/issues/2985.
-@@ -72,10 +70,7 @@
+@@ -72,10 +72,7 @@
  
  
  class Factory(t.Protocol):
@@ -136,7 +127,7 @@ elif   unformatted:
  
  
  # Regression test for https://github.com/psf/black/issues/3436.
-@@ -83,5 +78,5 @@
+@@ -83,5 +80,5 @@
      return x
  # fmt: off
  elif   unformatted:
@@ -160,6 +151,8 @@ setup(
             ]  # Includes an formatted indentation.
     },
 )
+
+
 # Regression test for https://github.com/psf/black/issues/2015.
 run(
     # fmt: off

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__function_trailing_comma_py.snap

@@ -188,11 +188,8 @@ some_module.some_function(
  ):
      pass
  
-@@ -100,15 +56,7 @@
- some_module.some_function(
-     argument1, (one_element_tuple,), argument4, argument5, argument6
- )
--
+@@ -103,12 +59,5 @@
+ 
  # Inner trailing comma causes outer to explode
  some_module.some_function(
 -    argument1,
@@ -268,6 +265,7 @@ def func() -> ((also_super_long_type_annotation_that_may_cause_an_AST_related_cr
 some_module.some_function(
     argument1, (one_element_tuple,), argument4, argument5, argument6
 )
+
 # Inner trailing comma causes outer to explode
 some_module.some_function(
     argument1, (one, two,), argument4, argument5, argument6

crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__import_spacing_py.snap

@@ -62,7 +62,7 @@ __all__ = (
 ```diff
 --- Black
 +++ Ruff
-@@ -2,12 +2,13 @@
+@@ -2,8 +2,10 @@
  
  # flake8: noqa
  
@@ -74,11 +74,7 @@ __all__ = (
      ERROR,
  )
  import sys
--
- # This relies on each of the submodules having an __all__ variable.
- from .base_events import *
- from .coroutines import *
-@@ -22,33 +23,16 @@
+@@ -22,33 +24,16 @@
  from ..streams import *
  
  from some_library import (
@@ -134,6 +130,7 @@ from logging import (
     ERROR,
 )
 import sys
+
 # This relies on each of the submodules having an __all__ variable.
 from .base_events import *
 from .coroutines import *