From bc63cc9b3c055b5091da1a8bf433cfb95f77f75c Mon Sep 17 00:00:00 2001
From: konstin <konstin@mailbox.org>
Date: Wed, 21 Jun 2023 19:45:53 +0200
Subject: [PATCH] Fix remaining CPython formatter errors except for function
 argument separator comments (#5210)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

## Summary

This fixes two problems discovered when trying to format the cpython
repo with `cargo run --bin ruff_dev -- check-formatter-stability
projects/cpython`:

The first is to ignore try/except trailing comments for now since they
lead to unstable formatting on the dummy.

The second is to avoid dropping trailing if comments through placement:
This changes the placement to keep a comment trailing an if-elif or
if-elif-else to keep the comment a trailing comment on the entire if.
Previously the last comment would have been lost.
```python
if "first if":
    pass
elif "first elif":
    pass
```

The last remaining problem in cpython so far is function signature
argument separator comment placement which is its own PR on top of this.

## Test Plan

I added test fixtures of minimized examples with links back to the
original cpython location
---
 .../test/fixtures/ruff/statement/if.py        | 27 ++++++++++
 .../src/comments/placement.rs                 | 24 ++++++++-
 ...atter__tests__black_test__comments_py.snap | 40 +++++++-------
 ...__black_test__remove_except_parens_py.snap | 32 ++++++-----
 ...r__tests__ruff_test__statement__if_py.snap | 54 +++++++++++++++++++
 .../src/statement/stmt_try.rs                 |  5 ++
 6 files changed, 148 insertions(+), 34 deletions(-)

diff --git a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/if.py b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/if.py
index 65b21e30e9..91303ef592 100644
--- a/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/if.py
+++ b/crates/ruff_python_formatter/resources/test/fixtures/ruff/statement/if.py
@@ -35,3 +35,30 @@ elif aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa +
 
 else:
     ...
+
+# Regression test: Don't drop the trailing comment by associating it with the elif
+# instead of the else.
+# Originally found in https://github.com/python/cpython/blob/ab3823a97bdeefb0266b3c8d493f7f6223ce3686/Lib/dataclasses.py#L539
+
+if "if 1":
+    pass
+elif "elif 1":
+    pass
+# Don't drop this comment 1
+x = 1
+
+if "if 2":
+    pass
+elif "elif 2":
+    pass
+else:
+    pass
+# Don't drop this comment 2
+x = 2
+
+if "if 3":
+    pass
+else:
+    pass
+# Don't drop this comment 3
+x = 3
diff --git a/crates/ruff_python_formatter/src/comments/placement.rs b/crates/ruff_python_formatter/src/comments/placement.rs
index fa6fa5baaa..32c795fee2 100644
--- a/crates/ruff_python_formatter/src/comments/placement.rs
+++ b/crates/ruff_python_formatter/src/comments/placement.rs
@@ -402,7 +402,11 @@ fn handle_trailing_body_comment<'a>(
     }
 
     // Only do something if the preceding node has a body (has indented statements).
-    let Some(last_child) = comment.preceding_node().and_then(last_child_in_body) else {
+    let Some(preceding_node) = comment.preceding_node() else {
+        return CommentPlacement::Default(comment);
+    };
+
+    let Some(last_child) = last_child_in_body(preceding_node) else {
         return CommentPlacement::Default(comment);
     };
 
@@ -415,8 +419,24 @@ fn handle_trailing_body_comment<'a>(
     // the indent-level doesn't depend on the tab width (the indent level must be the same if the tab width is 1 or 8).
     let comment_indentation_len = comment_indentation.len();
 
+    // Keep the comment on the entire statement in case it's a trailing comment
+    // ```python
+    // if "first if":
+    //     pass
+    // elif "first elif":
+    //     pass
+    // # Trailing if comment
+    // ```
+    // Here we keep the comment a trailing comment of the `if`
+    let Some(preceding_node_indentation) = whitespace::indentation_at_offset(locator, preceding_node.start()) else {
+        return CommentPlacement::Default(comment);
+    };
+    if comment_indentation_len == preceding_node_indentation.len() {
+        return CommentPlacement::Default(comment);
+    }
+
     let mut current_child = last_child;
-    let mut parent_body = comment.preceding_node();
+    let mut parent_body = Some(preceding_node);
     let mut grand_parent_body = None;
 
     loop {
diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments_py.snap
index 5a2dcb790f..e9ed496189 100644
--- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments_py.snap
+++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__comments_py.snap
@@ -109,38 +109,35 @@ async def wat():
 ```diff
 --- Black
 +++ Ruff
-@@ -4,24 +4,15 @@
+@@ -4,21 +4,15 @@
  #
  # Has many lines. Many, many lines.
  # Many, many, many lines.
 -"""Module docstring.
-+"NOT_YET_IMPLEMENTED_STRING"
- 
+-
 -Possibly also many, many lines.
 -"""
-+NOT_YET_IMPLEMENTED_StmtImport
-+NOT_YET_IMPLEMENTED_StmtImport
++"NOT_YET_IMPLEMENTED_STRING"
  
 -import os.path
 -import sys
--
++NOT_YET_IMPLEMENTED_StmtImport
++NOT_YET_IMPLEMENTED_StmtImport
+ 
 -import a
 -from b.c import X  # some noqa comment
--
++NOT_YET_IMPLEMENTED_StmtImport
++NOT_YET_IMPLEMENTED_StmtImportFrom  # some noqa comment
+ 
 -try:
 -    import fast
 -except ImportError:
 -    import slow as fast
--
-+NOT_YET_IMPLEMENTED_StmtImport
-+NOT_YET_IMPLEMENTED_StmtImportFrom  # some noqa comment
- 
--# Some comment before a function.
 +NOT_YET_IMPLEMENTED_StmtTry
- y = 1
- (
-     # some strings
-@@ -30,67 +21,50 @@
+ 
+ 
+ # Some comment before a function.
+@@ -30,67 +24,50 @@
  
  
  def function(default=None):
@@ -177,17 +174,17 @@ async def wat():
  
  # Another comment!
  # This time two lines.
--
--
+ 
+ 
 -class Foo:
 -    """Docstring for class Foo.  Example from Sphinx docs."""
- 
+-
 -    #: Doc comment for class attribute Foo.bar.
 -    #: It can have multiple lines.
 -    bar = 1
 -
 -    flox = 1.5  #: Doc comment for Foo.flox. One line only.
- 
+-
 -    baz = 2
 -    """Docstring for class attribute Foo.baz."""
 -
@@ -245,6 +242,9 @@ NOT_YET_IMPLEMENTED_StmtImport
 NOT_YET_IMPLEMENTED_StmtImportFrom  # some noqa comment
 
 NOT_YET_IMPLEMENTED_StmtTry
+
+
+# Some comment before a function.
 y = 1
 (
     # some strings
diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_except_parens_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_except_parens_py.snap
index 21fdd1942d..680584a6d7 100644
--- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_except_parens_py.snap
+++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__black_test__remove_except_parens_py.snap
@@ -48,28 +48,31 @@ except (some.really.really.really.looooooooooooooooooooooooooooooooong.module.ov
 ```diff
 --- Black
 +++ Ruff
-@@ -1,42 +1,9 @@
+@@ -1,42 +1,17 @@
  # These brackets are redundant, therefore remove.
 -try:
 -    a.something
 -except AttributeError as err:
 -    raise err
--
--# This is tuple of exceptions.
--# Although this could be replaced with just the exception,
--# we do not remove brackets to preserve AST.
++NOT_YET_IMPLEMENTED_StmtTry
+ 
+ # This is tuple of exceptions.
+ # Although this could be replaced with just the exception,
+ # we do not remove brackets to preserve AST.
 -try:
 -    a.something
 -except (AttributeError,) as err:
 -    raise err
--
--# This is a tuple of exceptions. Do not remove brackets.
++NOT_YET_IMPLEMENTED_StmtTry
+ 
+ # This is a tuple of exceptions. Do not remove brackets.
 -try:
 -    a.something
 -except (AttributeError, ValueError) as err:
 -    raise err
--
--# Test long variants.
++NOT_YET_IMPLEMENTED_StmtTry
+ 
+ # Test long variants.
 -try:
 -    a.something
 -except (
@@ -77,9 +80,6 @@ except (some.really.really.really.looooooooooooooooooooooooooooooooong.module.ov
 -) as err:
 -    raise err
 +NOT_YET_IMPLEMENTED_StmtTry
-+NOT_YET_IMPLEMENTED_StmtTry
-+NOT_YET_IMPLEMENTED_StmtTry
-+NOT_YET_IMPLEMENTED_StmtTry
  
 -try:
 -    a.something
@@ -104,8 +104,16 @@ except (some.really.really.really.looooooooooooooooooooooooooooooooong.module.ov
 ```py
 # These brackets are redundant, therefore remove.
 NOT_YET_IMPLEMENTED_StmtTry
+
+# This is tuple of exceptions.
+# Although this could be replaced with just the exception,
+# we do not remove brackets to preserve AST.
 NOT_YET_IMPLEMENTED_StmtTry
+
+# This is a tuple of exceptions. Do not remove brackets.
 NOT_YET_IMPLEMENTED_StmtTry
+
+# Test long variants.
 NOT_YET_IMPLEMENTED_StmtTry
 
 NOT_YET_IMPLEMENTED_StmtTry
diff --git a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__statement__if_py.snap b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__statement__if_py.snap
index a58a992e81..ef06f0e165 100644
--- a/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__statement__if_py.snap
+++ b/crates/ruff_python_formatter/src/snapshots/ruff_python_formatter__tests__ruff_test__statement__if_py.snap
@@ -41,6 +41,33 @@ elif aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa +
 
 else:
     ...
+
+# Regression test: Don't drop the trailing comment by associating it with the elif
+# instead of the else.
+# Originally found in https://github.com/python/cpython/blob/ab3823a97bdeefb0266b3c8d493f7f6223ce3686/Lib/dataclasses.py#L539
+
+if "if 1":
+    pass
+elif "elif 1":
+    pass
+# Don't drop this comment 1
+x = 1
+
+if "if 2":
+    pass
+elif "elif 2":
+    pass
+else:
+    pass
+# Don't drop this comment 2
+x = 2
+
+if "if 3":
+    pass
+else:
+    pass
+# Don't drop this comment 3
+x = 3
 ```
 
 
@@ -83,6 +110,33 @@ elif aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa +
 
 else:
     ...
+
+# Regression test: Don't drop the trailing comment by associating it with the elif
+# instead of the else.
+# Originally found in https://github.com/python/cpython/blob/ab3823a97bdeefb0266b3c8d493f7f6223ce3686/Lib/dataclasses.py#L539
+
+if "NOT_YET_IMPLEMENTED_STRING":
+    pass
+elif "NOT_YET_IMPLEMENTED_STRING":
+    pass
+# Don't drop this comment 1
+x = 1
+
+if "NOT_YET_IMPLEMENTED_STRING":
+    pass
+elif "NOT_YET_IMPLEMENTED_STRING":
+    pass
+else:
+    pass
+# Don't drop this comment 2
+x = 2
+
+if "NOT_YET_IMPLEMENTED_STRING":
+    pass
+else:
+    pass
+# Don't drop this comment 3
+x = 3
 ```
 
 
diff --git a/crates/ruff_python_formatter/src/statement/stmt_try.rs b/crates/ruff_python_formatter/src/statement/stmt_try.rs
index 35afc06669..faab829230 100644
--- a/crates/ruff_python_formatter/src/statement/stmt_try.rs
+++ b/crates/ruff_python_formatter/src/statement/stmt_try.rs
@@ -9,4 +9,9 @@ impl FormatNodeRule<StmtTry> for FormatStmtTry {
     fn fmt_fields(&self, item: &StmtTry, f: &mut PyFormatter) -> FormatResult<()> {
         write!(f, [not_yet_implemented(item)])
     }
+
+    fn fmt_dangling_comments(&self, _node: &StmtTry, _f: &mut PyFormatter) -> FormatResult<()> {
+        // TODO(konstin): Needs node formatting or this leads to unstable formatting
+        Ok(())
+    }
 }
