Implement autofix for newline-related docstring rules (#441)

src/checks.rs

@@ -1205,6 +1205,12 @@ impl CheckKind {
             CheckKind::DeprecatedUnittestAlias(_, _)
                 | CheckKind::DoNotAssertFalse
                 | CheckKind::DuplicateHandlerException(_)
+                | CheckKind::NoBlankLineAfterFunction(_)
+                | CheckKind::NoBlankLineAfterSummary
+                | CheckKind::NoBlankLineBeforeClass(_)
+                | CheckKind::NoBlankLineBeforeFunction(_)
+                | CheckKind::OneBlankLineAfterClass(_)
+                | CheckKind::OneBlankLineBeforeClass(_)
                 | CheckKind::PPrintFound
                 | CheckKind::PrintFound
                 | CheckKind::SuperCallWithParameters
@@ -1212,10 +1218,10 @@ impl CheckKind {
                 | CheckKind::UnnecessaryAbspath
                 | CheckKind::UnusedImport(_)
                 | CheckKind::UnusedNOQA(_)
-                | CheckKind::UselessMetaclassType
-                | CheckKind::UselessObjectInheritance(_)
                 | CheckKind::UsePEP585Annotation(_)
                 | CheckKind::UsePEP604Annotation
+                | CheckKind::UselessMetaclassType
+                | CheckKind::UselessObjectInheritance(_)
         )
     }
 }
@@ -1228,6 +1234,26 @@ pub struct Fix {
     pub applied: bool,
 }
 
+impl Fix {
+    pub fn deletion(start: Location, end: Location) -> Self {
+        Self {
+            content: "".to_string(),
+            location: start,
+            end_location: end,
+            applied: false,
+        }
+    }
+
+    pub fn insertion(content: String, start: Location, end: Location) -> Self {
+        Self {
+            content,
+            location: start,
+            end_location: end,
+            applied: false,
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct Check {
     pub kind: CheckKind,
@@ -1237,11 +1263,11 @@ pub struct Check {
 }
 
 impl Check {
-    pub fn new(kind: CheckKind, span: Range) -> Self {
+    pub fn new(kind: CheckKind, rage: Range) -> Self {
         Self {
             kind,
-            location: span.location,
+            location: rage.location,
-            end_location: span.end_location,
+            end_location: rage.end_location,
             fix: None,
         }
     }

src/docstrings/plugins.rs

@@ -5,8 +5,9 @@ use regex::Regex;
 use rustpython_ast::{Constant, ExprKind, Location, StmtKind};
 
 use crate::ast::types::Range;
+use crate::autofix::fixer;
 use crate::check_ast::Checker;
-use crate::checks::{Check, CheckCode, CheckKind};
+use crate::checks::{Check, CheckCode, CheckKind, Fix};
 use crate::docstrings::google::check_google_section;
 use crate::docstrings::helpers::{indentation, leading_space};
 use crate::docstrings::numpy::check_numpy_section;
@@ -174,35 +175,57 @@ pub fn blank_before_after_function(checker: &mut Checker, definition: &Definitio
                         .take_while(|line| line.trim().is_empty())
                         .count();
                     if blank_lines_before != 0 {
-                        checker.add_check(Check::new(
+                        let mut check = Check::new(
                             CheckKind::NoBlankLineBeforeFunction(blank_lines_before),
                             Range::from_located(docstring),
-                        ));
+                        );
+                        if matches!(checker.autofix, fixer::Mode::Generate | fixer::Mode::Apply) {
+                            check.amend(Fix::deletion(
+                                Location::new(docstring.location.row() - blank_lines_before, 1),
+                                Location::new(docstring.location.row(), 1),
+                            ));
+                        }
+                        checker.add_check(check);
                     }
                 }
 
                 if checker.settings.enabled.contains(&CheckCode::D202) {
+                    let all_blank_after = after
+                        .lines()
+                        .skip(1)
+                        .all(|line| line.trim().is_empty() || COMMENT_REGEX.is_match(line));
+                    if all_blank_after {
+                        return;
+                    }
+
                     let blank_lines_after = after
                         .lines()
                         .skip(1)
                         .take_while(|line| line.trim().is_empty())
                         .count();
-                    let all_blank_after = after
+                    // Report a D202 violation if the docstring is followed by a blank line and the
-                        .lines()
+                    // blank line is not itself followed by an inner function or class.
-                        .skip(1)
+                    let expected_blank_lines_after =
-                        .all(|line| line.trim().is_empty() || COMMENT_REGEX.is_match(line));
+                        if INNER_FUNCTION_OR_CLASS_REGEX.is_match(after) {
-                    // Report a D202 violation if the docstring is followed by a blank line
+                            1
-                    // and the blank line is not itself followed by an inner function or
+                        } else {
-                    // class.
+                            0
-                    if !all_blank_after
+                        };
-                        && blank_lines_after != 0
+                    if blank_lines_after != expected_blank_lines_after {
-                        && !(blank_lines_after == 1
+                        let mut check = Check::new(
-                            && INNER_FUNCTION_OR_CLASS_REGEX.is_match(after))
-                    {
-                        checker.add_check(Check::new(
                             CheckKind::NoBlankLineAfterFunction(blank_lines_after),
                             Range::from_located(docstring),
-                        ));
+                        );
+                        if matches!(checker.autofix, fixer::Mode::Generate | fixer::Mode::Apply) {
+                            check.amend(Fix::deletion(
+                                Location::new(
+                                    docstring.location.row() + 1 + expected_blank_lines_after,
+                                    1,
+                                ),
+                                Location::new(docstring.location.row() + 1 + blank_lines_after, 1),
+                            ));
+                        }
+                        checker.add_check(check);
                     }
                 }
             }
@@ -235,39 +258,71 @@ pub fn blank_before_after_class(checker: &mut Checker, definition: &Definition)
                         .skip(1)
                         .take_while(|line| line.trim().is_empty())
                         .count();
-                    if blank_lines_before != 0
+                    if checker.settings.enabled.contains(&CheckCode::D211) {
-                        && checker.settings.enabled.contains(&CheckCode::D211)
+                        if blank_lines_before != 0 {
-                    {
+                            let mut check = Check::new(
-                        checker.add_check(Check::new(
+                                CheckKind::NoBlankLineBeforeClass(blank_lines_before),
-                            CheckKind::NoBlankLineBeforeClass(blank_lines_before),
+                                Range::from_located(docstring),
-                            Range::from_located(docstring),
+                            );
-                        ));
+                            if matches!(checker.autofix, fixer::Mode::Generate | fixer::Mode::Apply)
+                            {
+                                check.amend(Fix::deletion(
+                                    Location::new(docstring.location.row() - blank_lines_before, 1),
+                                    Location::new(docstring.location.row(), 1),
+                                ));
+                            }
+                            checker.add_check(check);
+                        }
                     }
-                    if blank_lines_before != 1
+                    if checker.settings.enabled.contains(&CheckCode::D203) {
-                        && checker.settings.enabled.contains(&CheckCode::D203)
+                        if blank_lines_before != 1 {
-                    {
+                            let mut check = Check::new(
-                        checker.add_check(Check::new(
+                                CheckKind::OneBlankLineBeforeClass(blank_lines_before),
-                            CheckKind::OneBlankLineBeforeClass(blank_lines_before),
+                                Range::from_located(docstring),
-                            Range::from_located(docstring),
+                            );
-                        ));
+                            if matches!(checker.autofix, fixer::Mode::Generate | fixer::Mode::Apply)
+                            {
+                                check.amend(Fix::insertion(
+                                    "\n".to_string(),
+                                    Location::new(docstring.location.row() - blank_lines_before, 1),
+                                    Location::new(docstring.location.row(), 1),
+                                ));
+                            }
+                            checker.add_check(check);
+                        }
                     }
                 }
 
                 if checker.settings.enabled.contains(&CheckCode::D204) {
+                    let all_blank_after = after
+                        .lines()
+                        .skip(1)
+                        .all(|line| line.trim().is_empty() || COMMENT_REGEX.is_match(line));
+                    if all_blank_after {
+                        return;
+                    }
+
                     let blank_lines_after = after
                         .lines()
                         .skip(1)
                         .take_while(|line| line.trim().is_empty())
                         .count();
-                    let all_blank_after = after
+                    if blank_lines_after != 1 {
-                        .lines()
+                        let mut check = Check::new(
-                        .skip(1)
-                        .all(|line| line.trim().is_empty() || COMMENT_REGEX.is_match(line));
-                    if !all_blank_after && blank_lines_after != 1 {
-                        checker.add_check(Check::new(
                             CheckKind::OneBlankLineAfterClass(blank_lines_after),
                             Range::from_located(docstring),
-                        ));
+                        );
+                        if matches!(checker.autofix, fixer::Mode::Generate | fixer::Mode::Apply) {
+                            check.amend(Fix::insertion(
+                                "\n".to_string(),
+                                Location::new(docstring.end_location.row() + 1, 1),
+                                Location::new(
+                                    docstring.end_location.row() + 1 + blank_lines_after,
+                                    1,
+                                ),
+                            ));
+                        }
+                        checker.add_check(check);
                     }
                 }
             }
@@ -294,10 +349,18 @@ pub fn blank_after_summary(checker: &mut Checker, definition: &Definition) {
                 }
             }
             if lines_count > 1 && blanks_count != 1 {
-                checker.add_check(Check::new(
+                let mut check = Check::new(
                     CheckKind::NoBlankLineAfterSummary,
                     Range::from_located(docstring),
-                ));
+                );
+                if matches!(checker.autofix, fixer::Mode::Generate | fixer::Mode::Apply) {
+                    check.amend(Fix::insertion(
+                        "\n".to_string(),
+                        Location::new(docstring.location.row() + 1, 1),
+                        Location::new(docstring.location.row() + 1 + blanks_count, 1),
+                    ));
+                }
+                checker.add_check(check);
             }
         }
     }

src/snapshots/ruff__linter__tests__D201_D.py.snap

@@ -10,7 +10,15 @@ expression: checks
   end_location:
     row: 132
     column: 25
-  fix: ~
+  fix:
+    content: ""
+    location:
+      row: 131
+      column: 1
+    end_location:
+      row: 132
+      column: 1
+    applied: false
 - kind:
     NoBlankLineBeforeFunction: 1
   location:
@@ -19,5 +27,13 @@ expression: checks
   end_location:
     row: 146
     column: 38
-  fix: ~
+  fix:
+    content: ""
+    location:
+      row: 145
+      column: 1
+    end_location:
+      row: 146
+      column: 1
+    applied: false
 

src/snapshots/ruff__linter__tests__D202_D.py.snap

@@ -2,6 +2,23 @@
 source: src/linter.rs
 expression: checks
 ---
+- kind:
+    NoBlankLineAfterFunction: 0
+  location:
+    row: 79
+    column: 5
+  end_location:
+    row: 79
+    column: 33
+  fix:
+    content: ""
+    location:
+      row: 81
+      column: 1
+    end_location:
+      row: 80
+      column: 1
+    applied: false
 - kind:
     NoBlankLineAfterFunction: 1
   location:
@@ -10,7 +27,15 @@ expression: checks
   end_location:
     row: 137
     column: 25
-  fix: ~
+  fix:
+    content: ""
+    location:
+      row: 138
+      column: 1
+    end_location:
+      row: 139
+      column: 1
+    applied: false
 - kind:
     NoBlankLineAfterFunction: 1
   location:
@@ -19,5 +44,30 @@ expression: checks
   end_location:
     row: 146
     column: 38
-  fix: ~
+  fix:
+    content: ""
+    location:
+      row: 147
+      column: 1
+    end_location:
+      row: 148
+      column: 1
+    applied: false
+- kind:
+    NoBlankLineAfterFunction: 0
+  location:
+    row: 453
+    column: 5
+  end_location:
+    row: 453
+    column: 24
+  fix:
+    content: ""
+    location:
+      row: 455
+      column: 1
+    end_location:
+      row: 454
+      column: 1
+    applied: false
 

src/snapshots/ruff__linter__tests__D203_D.py.snap

@@ -10,7 +10,15 @@ expression: checks
   end_location:
     row: 156
     column: 33
-  fix: ~
+  fix:
+    content: "\n"
+    location:
+      row: 156
+      column: 1
+    end_location:
+      row: 156
+      column: 1
+    applied: false
 - kind:
     OneBlankLineBeforeClass: 0
   location:
@@ -19,7 +27,15 @@ expression: checks
   end_location:
     row: 187
     column: 46
-  fix: ~
+  fix:
+    content: "\n"
+    location:
+      row: 187
+      column: 1
+    end_location:
+      row: 187
+      column: 1
+    applied: false
 - kind:
     OneBlankLineBeforeClass: 0
   location:
@@ -28,5 +44,13 @@ expression: checks
   end_location:
     row: 527
     column: 8
-  fix: ~
+  fix:
+    content: "\n"
+    location:
+      row: 521
+      column: 1
+    end_location:
+      row: 521
+      column: 1
+    applied: false
 

src/snapshots/ruff__linter__tests__D204_D.py.snap

@@ -10,7 +10,15 @@ expression: checks
   end_location:
     row: 176
     column: 25
-  fix: ~
+  fix:
+    content: "\n"
+    location:
+      row: 177
+      column: 1
+    end_location:
+      row: 177
+      column: 1
+    applied: false
 - kind:
     OneBlankLineAfterClass: 0
   location:
@@ -19,5 +27,13 @@ expression: checks
   end_location:
     row: 187
     column: 46
-  fix: ~
+  fix:
+    content: "\n"
+    location:
+      row: 188
+      column: 1
+    end_location:
+      row: 188
+      column: 1
+    applied: false
 

src/snapshots/ruff__linter__tests__D205_D.py.snap

@@ -9,7 +9,15 @@ expression: checks
   end_location:
     row: 198
     column: 8
-  fix: ~
+  fix:
+    content: "\n"
+    location:
+      row: 196
+      column: 1
+    end_location:
+      row: 196
+      column: 1
+    applied: false
 - kind: NoBlankLineAfterSummary
   location:
     row: 205
@@ -17,5 +25,13 @@ expression: checks
   end_location:
     row: 210
     column: 8
-  fix: ~
+  fix:
+    content: "\n"
+    location:
+      row: 206
+      column: 1
+    end_location:
+      row: 208
+      column: 1
+    applied: false
 

src/snapshots/ruff__linter__tests__D211_D.py.snap

@@ -10,7 +10,15 @@ expression: checks
   end_location:
     row: 165
     column: 30
-  fix: ~
+  fix:
+    content: ""
+    location:
+      row: 164
+      column: 1
+    end_location:
+      row: 165
+      column: 1
+    applied: false
 - kind:
     NoBlankLineBeforeClass: 1
   location:
@@ -19,5 +27,13 @@ expression: checks
   end_location:
     row: 176
     column: 25
-  fix: ~
+  fix:
+    content: ""
+    location:
+      row: 175
+      column: 1
+    end_location:
+      row: 176
+      column: 1
+    applied: false