Fix D204 when there's a semicolon after a docstring (#7174)

## Summary

Another statement on the same line as the docstring would previous make
the D204 (newline after docstring) fix fail:

```python
class StatementOnSameLineAsDocstring:
    "After this docstring there's another statement on the same line separated by a semicolon." ;priorities=1
    def sort_services(self):
        pass
```

The fix handles this case manually:

```python
class StatementOnSameLineAsDocstring:
    "After this docstring there's another statement on the same line separated by a semicolon."

    priorities=1
    def sort_services(self):
        pass
```

Fixes #7088

## Test Plan

Added a new `D` test case

Author: konstin

crates/ruff/resources/test/fixtures/pydocstyle/D.py

@@ -644,3 +644,17 @@ def same_line(): """This is a docstring on the same line"""
 def single_line_docstring_with_an_escaped_backslash():
     "\
     "
+
+class StatementOnSameLineAsDocstring:
+    "After this docstring there's another statement on the same line separated by a semicolon." ; priorities=1
+    def sort_services(self):
+        pass
+
+class StatementOnSameLineAsDocstring:
+    "After this docstring there's another statement on the same line separated by a semicolon."; priorities=1
+
+
+class CommentAfterDocstring:
+    "After this docstring there's a comment."  # priorities=1
+    def sort_services(self):
+        pass

crates/ruff/src/rules/pydocstyle/rules/blank_before_after_class.rs

@@ -1,7 +1,7 @@
 use ruff_diagnostics::{AlwaysAutofixableViolation, Diagnostic, Edit, Fix};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_trivia::PythonWhitespace;
-use ruff_source_file::{UniversalNewlineIterator, UniversalNewlines};
+use ruff_python_trivia::{indentation_at_offset, PythonWhitespace};
+use ruff_source_file::{Line, UniversalNewlineIterator};
 use ruff_text_size::Ranged;
 use ruff_text_size::{TextLen, TextRange};
 
@@ -216,25 +216,61 @@ pub(crate) fn blank_before_after_class(checker: &mut Checker, docstring: &Docstr
     }
 
     if checker.enabled(Rule::OneBlankLineAfterClass) {
-        let after_range = TextRange::new(docstring.end(), class.end());
-
-        let after = checker.locator().slice(after_range);
+        let class_after_docstring_range = TextRange::new(docstring.end(), class.end());
+        let class_after_docstring = checker.locator().slice(class_after_docstring_range);
+        let mut lines = UniversalNewlineIterator::with_offset(
+            class_after_docstring,
+            class_after_docstring_range.start(),
+        );
 
-        let all_blank_after = after.universal_newlines().skip(1).all(|line| {
+        // If the class is empty except for comments, we don't need to insert a newline between
+        // docstring and no content
+        let all_blank_after = lines.clone().all(|line| {
             line.trim_whitespace().is_empty() || line.trim_whitespace_start().starts_with('#')
         });
         if all_blank_after {
             return;
         }
 
-        let mut lines = UniversalNewlineIterator::with_offset(after, after_range.start());
+        let first_line = lines.next();
+        let mut replacement_start = first_line.as_ref().map(Line::start).unwrap_or_default();
+
+        // Edge case: There is trailing end-of-line content after the docstring, either a statement
+        // separated by a semicolon or a comment.
+        if let Some(first_line) = &first_line {
+            let trailing = first_line.as_str().trim_whitespace_start();
+            if let Some(next_statement) = trailing.strip_prefix(';') {
+                let indentation = indentation_at_offset(docstring.start(), checker.locator())
+                    .expect("Own line docstring must have indentation");
+                let mut diagnostic = Diagnostic::new(OneBlankLineAfterClass, docstring.range());
+                if checker.patch(diagnostic.kind.rule()) {
+                    let line_ending = checker.stylist().line_ending().as_str();
+                    // We have to trim the whitespace twice, once before the semicolon above and
+                    // once after the semicolon here, or we get invalid indents:
+                    // ```rust
+                    // class Priority:
+                    //     """Has priorities"""   ;   priorities=1
+                    // ```
+                    let next_statement = next_statement.trim_whitespace_start();
+                    diagnostic.set_fix(Fix::automatic(Edit::replacement(
+                        line_ending.to_string() + line_ending + indentation + next_statement,
+                        replacement_start,
+                        first_line.end(),
+                    )));
+                }
+                checker.diagnostics.push(diagnostic);
+                return;
+            } else if trailing.starts_with('#') {
+                // Keep the end-of-line comment, start counting empty lines after it
+                replacement_start = first_line.end();
+            }
+        }
 
-        let first_line_start = lines.next().map(|l| l.start()).unwrap_or_default();
         let mut blank_lines_after = 0usize;
-        let mut blank_lines_end = docstring.end();
+        let mut blank_lines_end = first_line.as_ref().map_or(docstring.end(), Line::end);
 
         for line in lines {
-            if line.trim().is_empty() {
+            if line.trim_whitespace().is_empty() {
                 blank_lines_end = line.end();
                 blank_lines_after += 1;
             } else {
@@ -248,7 +284,7 @@ pub(crate) fn blank_before_after_class(checker: &mut Checker, docstring: &Docstr
                 // Insert a blank line before the class (replacing any existing lines).
                 diagnostic.set_fix(Fix::automatic(Edit::replacement(
                     checker.stylist().line_ending().to_string(),
-                    first_line_start,
+                    replacement_start,
                     blank_lines_end,
                 )));
             }

crates/ruff/src/rules/pydocstyle/snapshots/ruff__rules__pydocstyle__tests__D102_D.py.snap

@@ -25,4 +25,22 @@ D.py:68:9: D102 Missing docstring in public method
 69 |         pass
    |
 
+D.py:650:9: D102 Missing docstring in public method
+    |
+648 | class StatementOnSameLineAsDocstring:
+649 |     "After this docstring there's another statement on the same line separated by a semicolon." ; priorities=1
+650 |     def sort_services(self):
+    |         ^^^^^^^^^^^^^ D102
+651 |         pass
+    |
+
+D.py:659:9: D102 Missing docstring in public method
+    |
+657 | class CommentAfterDocstring:
+658 |     "After this docstring there's a comment."  # priorities=1
+659 |     def sort_services(self):
+    |         ^^^^^^^^^^^^^ D102
+660 |         pass
+    |
+
 

crates/ruff/src/rules/pydocstyle/snapshots/ruff__rules__pydocstyle__tests__D200_D.py.snap

@@ -104,6 +104,8 @@ D.py:645:5: D200 One-line docstring should fit on one line
     |  _____^
 646 | |     "
     | |_____^ D200
+647 |   
+648 |   class StatementOnSameLineAsDocstring:
     |
     = help: Reformat to one line
 

crates/ruff/src/rules/pydocstyle/snapshots/ruff__rules__pydocstyle__tests__D203_D.py.snap

@@ -63,4 +63,59 @@ D.py:526:5: D203 [*] 1 blank line required before class docstring
 527 528 | 
 528 529 |     Parameters
 
+D.py:649:5: D203 [*] 1 blank line required before class docstring
+    |
+648 | class StatementOnSameLineAsDocstring:
+649 |     "After this docstring there's another statement on the same line separated by a semicolon." ; priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D203
+650 |     def sort_services(self):
+651 |         pass
+    |
+    = help: Insert 1 blank line before class docstring
+
+ℹ Fix
+646 646 |     "
+647 647 | 
+648 648 | class StatementOnSameLineAsDocstring:
+    649 |+
+649 650 |     "After this docstring there's another statement on the same line separated by a semicolon." ; priorities=1
+650 651 |     def sort_services(self):
+651 652 |         pass
+
+D.py:654:5: D203 [*] 1 blank line required before class docstring
+    |
+653 | class StatementOnSameLineAsDocstring:
+654 |     "After this docstring there's another statement on the same line separated by a semicolon."; priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D203
+    |
+    = help: Insert 1 blank line before class docstring
+
+ℹ Fix
+651 651 |         pass
+652 652 | 
+653 653 | class StatementOnSameLineAsDocstring:
+    654 |+
+654 655 |     "After this docstring there's another statement on the same line separated by a semicolon."; priorities=1
+655 656 | 
+656 657 | 
+
+D.py:658:5: D203 [*] 1 blank line required before class docstring
+    |
+657 | class CommentAfterDocstring:
+658 |     "After this docstring there's a comment."  # priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D203
+659 |     def sort_services(self):
+660 |         pass
+    |
+    = help: Insert 1 blank line before class docstring
+
+ℹ Fix
+655 655 | 
+656 656 | 
+657 657 | class CommentAfterDocstring:
+    658 |+
+658 659 |     "After this docstring there's a comment."  # priorities=1
+659 660 |     def sort_services(self):
+660 661 |         pass
+
 

crates/ruff/src/rules/pydocstyle/snapshots/ruff__rules__pydocstyle__tests__D204_D.py.snap

@@ -38,4 +38,64 @@ D.py:192:5: D204 [*] 1 blank line required after class docstring
 194 195 | 
 195 196 | 
 
+D.py:649:5: D204 [*] 1 blank line required after class docstring
+    |
+648 | class StatementOnSameLineAsDocstring:
+649 |     "After this docstring there's another statement on the same line separated by a semicolon." ; priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D204
+650 |     def sort_services(self):
+651 |         pass
+    |
+    = help: Insert 1 blank line after class docstring
+
+ℹ Fix
+646 646 |     "
+647 647 | 
+648 648 | class StatementOnSameLineAsDocstring:
+649     |-    "After this docstring there's another statement on the same line separated by a semicolon." ; priorities=1
+    649 |+    "After this docstring there's another statement on the same line separated by a semicolon."
+    650 |+
+    651 |+    priorities=1
+650 652 |     def sort_services(self):
+651 653 |         pass
+652 654 | 
+
+D.py:654:5: D204 [*] 1 blank line required after class docstring
+    |
+653 | class StatementOnSameLineAsDocstring:
+654 |     "After this docstring there's another statement on the same line separated by a semicolon."; priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D204
+    |
+    = help: Insert 1 blank line after class docstring
+
+ℹ Fix
+651 651 |         pass
+652 652 | 
+653 653 | class StatementOnSameLineAsDocstring:
+654     |-    "After this docstring there's another statement on the same line separated by a semicolon."; priorities=1
+    654 |+    "After this docstring there's another statement on the same line separated by a semicolon."
+    655 |+
+    656 |+    priorities=1
+655 657 | 
+656 658 | 
+657 659 | class CommentAfterDocstring:
+
+D.py:658:5: D204 [*] 1 blank line required after class docstring
+    |
+657 | class CommentAfterDocstring:
+658 |     "After this docstring there's a comment."  # priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D204
+659 |     def sort_services(self):
+660 |         pass
+    |
+    = help: Insert 1 blank line after class docstring
+
+ℹ Fix
+656 656 | 
+657 657 | class CommentAfterDocstring:
+658 658 |     "After this docstring there's a comment."  # priorities=1
+    659 |+
+659 660 |     def sort_services(self):
+660 661 |         pass
+
 

crates/ruff/src/rules/pydocstyle/snapshots/ruff__rules__pydocstyle__tests__D300_D.py.snap

@@ -48,6 +48,33 @@ D.py:645:5: D300 Use triple double quotes `"""`
     |  _____^
 646 | |     "
     | |_____^ D300
+647 |   
+648 |   class StatementOnSameLineAsDocstring:
+    |
+
+D.py:649:5: D300 Use triple double quotes `"""`
+    |
+648 | class StatementOnSameLineAsDocstring:
+649 |     "After this docstring there's another statement on the same line separated by a semicolon." ; priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D300
+650 |     def sort_services(self):
+651 |         pass
+    |
+
+D.py:654:5: D300 Use triple double quotes `"""`
+    |
+653 | class StatementOnSameLineAsDocstring:
+654 |     "After this docstring there's another statement on the same line separated by a semicolon."; priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D300
+    |
+
+D.py:658:5: D300 Use triple double quotes `"""`
+    |
+657 | class CommentAfterDocstring:
+658 |     "After this docstring there's a comment."  # priorities=1
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ D300
+659 |     def sort_services(self):
+660 |         pass
     |
 
 

crates/ruff_source_file/src/newlines.rs

@@ -32,6 +32,7 @@ impl UniversalNewlines for str {
 /// assert_eq!(lines.next_back(), Some(Line::new("\r\n", TextSize::from(8))));
 /// assert_eq!(lines.next(), None);
 /// ```
+#[derive(Clone)]
 pub struct UniversalNewlineIterator<'a> {
     text: &'a str,
     offset: TextSize,