Implement D201, D202, D203, D204, and D211 (#404)

README.md

@@ -217,7 +217,7 @@ ruff also implements some of the most popular Flake8 plugins natively, including
 - [`flake8-print`](https://pypi.org/project/flake8-print/)
 - [`flake8-comprehensions`](https://pypi.org/project/flake8-comprehensions/) (11/16)
 - [`flake8-bugbear`](https://pypi.org/project/flake8-bugbear/) (3/32)
-- [`flake8-docstrings`](https://pypi.org/project/flake8-docstrings/) (12/47)
+- [`flake8-docstrings`](https://pypi.org/project/flake8-docstrings/) (17/48)
 - [`pyupgrade`](https://pypi.org/project/pyupgrade/) (8/34)
 
 Beyond rule-set parity, ruff suffers from the following limitations vis-à-vis Flake8:
@@ -316,6 +316,11 @@ The 🛠 emoji indicates that a rule is automatically fixable by the `--fix` com
 | D403 | FirstLineCapitalized | First word of the first line should be properly capitalized |  |  |
 | D415 | EndsInPunctuation | First line should end with a period, question mark, or exclamation point |  |  |
 | D419 | NonEmpty | Docstring is empty |  |  |
+| D201 | NoBlankLineBeforeFunction | No blank lines allowed before function docstring (found 1) |  |  |
+| D202 | NoBlankLineAfterFunction | No blank lines allowed after function docstring (found 1) |  |  |
+| D211 | NoBlankLineBeforeClass | NoBlankLineBeforeClass |  |  |
+| D203 | OneBlankLineBeforeClass | OneBlankLineBeforeClass |  |  |
+| D204 | OneBlankLineAfterClass | OneBlankLineAfterClass |  |  |
 | M001 | UnusedNOQA | Unused `noqa` directive |  | 🛠 |
 
 ## Integrations

resources/test/fixtures/D.py

@@ -1,4 +1,4 @@
-# No docstring, so we can test D100
+r# No docstring, so we can test D100
 from functools import wraps
 import os
 from .expected import Expectation

src/ast/operations.rs

@@ -134,7 +134,7 @@ impl<'a> SourceCodeLocator<'a> {
         }
     }
 
-    pub fn slice_source_code_at(&mut self, location: &Location) -> &'a str {
+    fn init(&mut self) {
         if !self.initialized {
             let mut offset = 0;
             for i in self.content.lines() {
@@ -142,24 +142,40 @@ impl<'a> SourceCodeLocator<'a> {
                 offset += i.len();
                 offset += 1;
             }
+            self.offsets.push(offset);
             self.initialized = true;
         }
+    }
+
+    pub fn slice_source_code_at(&mut self, location: &Location) -> &'a str {
+        self.init();
         let offset = self.offsets[location.row() - 1] + location.column() - 1;
         &self.content[offset..]
     }
 
     pub fn slice_source_code_range(&mut self, range: &Range) -> &'a str {
-        if !self.initialized {
-            let mut offset = 0;
-            for i in self.content.lines() {
-                self.offsets.push(offset);
-                offset += i.len();
-                offset += 1;
-            }
-            self.initialized = true;
-        }
+        self.init();
         let start = self.offsets[range.location.row() - 1] + range.location.column() - 1;
         let end = self.offsets[range.end_location.row() - 1] + range.end_location.column() - 1;
         &self.content[start..end]
     }
+
+    pub fn partition_source_code_at(
+        &mut self,
+        outer: &Range,
+        inner: &Range,
+    ) -> (&'a str, &'a str, &'a str) {
+        self.init();
+        let outer_start = self.offsets[outer.location.row() - 1] + outer.location.column() - 1;
+        let outer_end =
+            self.offsets[outer.end_location.row() - 1] + outer.end_location.column() - 1;
+        let inner_start = self.offsets[inner.location.row() - 1] + inner.location.column() - 1;
+        let inner_end =
+            self.offsets[inner.end_location.row() - 1] + inner.end_location.column() - 1;
+        (
+            &self.content[outer_start..inner_start],
+            &self.content[inner_start..inner_end],
+            &self.content[inner_end..outer_end],
+        )
+    }
 }

src/check_ast.rs

@@ -1896,6 +1896,17 @@ impl<'a> Checker<'a> {
             if self.settings.enabled.contains(&CheckCode::D200) {
                 docstrings::one_liner(self, &docstring);
             }
+            if self.settings.enabled.contains(&CheckCode::D201)
+                || self.settings.enabled.contains(&CheckCode::D202)
+            {
+                docstrings::blank_before_after_function(self, &docstring);
+            }
+            if self.settings.enabled.contains(&CheckCode::D203)
+                || self.settings.enabled.contains(&CheckCode::D204)
+                || self.settings.enabled.contains(&CheckCode::D211)
+            {
+                docstrings::blank_before_after_class(self, &docstring);
+            }
             if self.settings.enabled.contains(&CheckCode::D205) {
                 docstrings::blank_after_summary(self, &docstring);
             }

src/checks.rs

@@ -163,6 +163,11 @@ pub enum CheckCode {
     D403,
     D415,
     D419,
+    D201,
+    D202,
+    D211,
+    D203,
+    D204,
     // Meta
     M001,
 }
@@ -272,6 +277,11 @@ pub enum CheckKind {
     NonEmpty,
     UsesTripleQuotes,
     NoSignature,
+    NoBlankLineBeforeFunction(usize),
+    NoBlankLineAfterFunction(usize),
+    NoBlankLineBeforeClass(usize),
+    OneBlankLineBeforeClass(usize),
+    OneBlankLineAfterClass(usize),
     // Meta
     UnusedNOQA(Option<Vec<String>>),
 }
@@ -393,6 +403,11 @@ impl CheckCode {
             CheckCode::D402 => CheckKind::NoSignature,
             CheckCode::D403 => CheckKind::FirstLineCapitalized,
             CheckCode::D415 => CheckKind::EndsInPunctuation,
+            CheckCode::D201 => CheckKind::NoBlankLineBeforeFunction(1),
+            CheckCode::D202 => CheckKind::NoBlankLineAfterFunction(1),
+            CheckCode::D211 => CheckKind::NoBlankLineBeforeClass(1),
+            CheckCode::D203 => CheckKind::OneBlankLineBeforeClass(0),
+            CheckCode::D204 => CheckKind::OneBlankLineAfterClass(0),
             // Meta
             CheckCode::M001 => CheckKind::UnusedNOQA(None),
         }
@@ -493,6 +508,11 @@ impl CheckKind {
             CheckKind::NoSignature => &CheckCode::D402,
             CheckKind::FirstLineCapitalized => &CheckCode::D403,
             CheckKind::EndsInPunctuation => &CheckCode::D415,
+            CheckKind::NoBlankLineBeforeFunction(_) => &CheckCode::D201,
+            CheckKind::NoBlankLineAfterFunction(_) => &CheckCode::D202,
+            CheckKind::NoBlankLineBeforeClass(_) => &CheckCode::D211,
+            CheckKind::OneBlankLineBeforeClass(_) => &CheckCode::D203,
+            CheckKind::OneBlankLineAfterClass(_) => &CheckCode::D204,
             // Meta
             CheckKind::UnusedNOQA(_) => &CheckCode::M001,
         }
@@ -766,6 +786,15 @@ impl CheckKind {
             CheckKind::NoSignature => {
                 "First line should not be the function's 'signature'".to_string()
             }
+            CheckKind::NoBlankLineBeforeFunction(num_lines) => {
+                format!("No blank lines allowed before function docstring (found {num_lines})")
+            }
+            CheckKind::NoBlankLineAfterFunction(num_lines) => {
+                format!("No blank lines allowed after function docstring (found {num_lines})")
+            }
+            CheckKind::NoBlankLineBeforeClass(_) => "NoBlankLineBeforeClass".to_string(),
+            CheckKind::OneBlankLineBeforeClass(_) => "OneBlankLineBeforeClass".to_string(),
+            CheckKind::OneBlankLineAfterClass(_) => "OneBlankLineAfterClass".to_string(),
             // Meta
             CheckKind::UnusedNOQA(codes) => match codes {
                 None => "Unused `noqa` directive".to_string(),

src/docstrings.rs

@@ -1,3 +1,5 @@
+use once_cell::sync::Lazy;
+use regex::Regex;
 use rustpython_ast::{Constant, Expr, ExprKind, Location, Stmt, StmtKind};
 
 use crate::ast::types::Range;
@@ -85,7 +87,7 @@ pub fn one_liner(checker: &mut Checker, docstring: &Docstring) {
                 non_empty_line_count += 1;
             }
             if non_empty_line_count > 1 {
-                return;
+                break;
             }
         }
 
@@ -95,6 +97,121 @@ pub fn one_liner(checker: &mut Checker, docstring: &Docstring) {
     }
 }
 
+static COMMENT_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r"^\s*#").unwrap());
+
+static INNER_FUNCTION_OR_CLASS_REGEX: Lazy<Regex> =
+    Lazy::new(|| Regex::new(r"^\s+(?:(?:class|def|async def)\s|@)").unwrap());
+
+/// D201, D202
+pub fn blank_before_after_function(checker: &mut Checker, docstring: &Docstring) {
+    if let DocstringKind::Function(parent) = &docstring.kind {
+        if let ExprKind::Constant {
+            value: Constant::Str(_),
+            ..
+        } = &docstring.expr.node
+        {
+            let (before, _, after) = checker
+                .locator
+                .partition_source_code_at(&Range::from_located(parent), &range_for(docstring));
+
+            if checker.settings.enabled.contains(&CheckCode::D201) {
+                let blank_lines_before = before
+                    .lines()
+                    .rev()
+                    .skip(1)
+                    .take_while(|line| line.trim().is_empty())
+                    .count();
+                if blank_lines_before != 0 {
+                    checker.add_check(Check::new(
+                        CheckKind::NoBlankLineBeforeFunction(blank_lines_before),
+                        range_for(docstring),
+                    ));
+                }
+            }
+
+            if checker.settings.enabled.contains(&CheckCode::D202) {
+                let blank_lines_after = after
+                    .lines()
+                    .skip(1)
+                    .take_while(|line| line.trim().is_empty())
+                    .count();
+                let all_blank_after = after
+                    .lines()
+                    .skip(1)
+                    .all(|line| line.trim().is_empty() || COMMENT_REGEX.is_match(line));
+                // Report a D202 violation if the docstring is followed by a blank line
+                // and the blank line is not itself followed by an inner function or
+                // class.
+                if !all_blank_after
+                    && blank_lines_after != 0
+                    && !(blank_lines_after == 1 && INNER_FUNCTION_OR_CLASS_REGEX.is_match(after))
+                {
+                    checker.add_check(Check::new(
+                        CheckKind::NoBlankLineAfterFunction(blank_lines_after),
+                        range_for(docstring),
+                    ));
+                }
+            }
+        }
+    }
+}
+
+/// D203, D204, D211
+pub fn blank_before_after_class(checker: &mut Checker, docstring: &Docstring) {
+    if let DocstringKind::Class(parent) = &docstring.kind {
+        if let ExprKind::Constant {
+            value: Constant::Str(_),
+            ..
+        } = &docstring.expr.node
+        {
+            let (before, _, after) = checker
+                .locator
+                .partition_source_code_at(&Range::from_located(parent), &range_for(docstring));
+
+            if checker.settings.enabled.contains(&CheckCode::D203)
+                || checker.settings.enabled.contains(&CheckCode::D211)
+            {
+                let blank_lines_before = before
+                    .lines()
+                    .rev()
+                    .skip(1)
+                    .take_while(|line| line.trim().is_empty())
+                    .count();
+                if blank_lines_before != 0 && checker.settings.enabled.contains(&CheckCode::D211) {
+                    checker.add_check(Check::new(
+                        CheckKind::NoBlankLineBeforeClass(blank_lines_before),
+                        range_for(docstring),
+                    ));
+                }
+                if blank_lines_before != 1 && checker.settings.enabled.contains(&CheckCode::D203) {
+                    checker.add_check(Check::new(
+                        CheckKind::OneBlankLineBeforeClass(blank_lines_before),
+                        range_for(docstring),
+                    ));
+                }
+            }
+
+            if checker.settings.enabled.contains(&CheckCode::D204) {
+                let blank_lines_after = after
+                    .lines()
+                    .skip(1)
+                    .take_while(|line| line.trim().is_empty())
+                    .count();
+                let all_blank_after = after
+                    .lines()
+                    .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_for(docstring),
+                    ));
+                }
+            }
+        }
+    }
+}
+
 /// D205
 pub fn blank_after_summary(checker: &mut Checker, docstring: &Docstring) {
     if let ExprKind::Constant {
@@ -187,15 +304,7 @@ pub fn multi_line_summary_start(checker: &mut Checker, docstring: &Docstring) {
                 .slice_source_code_range(&range_for(docstring));
             if let Some(first_line) = content.lines().next() {
                 let first_line = first_line.trim();
-                if first_line == "\"\"\""
-                    || first_line == "'''"
-                    || first_line == "u\"\"\""
-                    || first_line == "u'''"
-                    || first_line == "r\"\"\""
-                    || first_line == "r'''"
-                    || first_line == "ur\"\"\""
-                    || first_line == "ur'''"
-                {
+                if first_line == "\"\"\"" || first_line == "'''" {
                     if checker.settings.enabled.contains(&CheckCode::D212) {
                         checker.add_check(Check::new(
                             CheckKind::MultiLineSummaryFirstLine,

src/linter.rs

@@ -1020,6 +1020,54 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn d201() -> Result<()> {
+        let mut checks = check_path(
+            Path::new("./resources/test/fixtures/D.py"),
+            &settings::Settings::for_rule(CheckCode::D201),
+            &fixer::Mode::Generate,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
+    #[test]
+    fn d202() -> Result<()> {
+        let mut checks = check_path(
+            Path::new("./resources/test/fixtures/D.py"),
+            &settings::Settings::for_rule(CheckCode::D202),
+            &fixer::Mode::Generate,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
+    #[test]
+    fn d203() -> Result<()> {
+        let mut checks = check_path(
+            Path::new("./resources/test/fixtures/D.py"),
+            &settings::Settings::for_rule(CheckCode::D203),
+            &fixer::Mode::Generate,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
+    #[test]
+    fn d204() -> Result<()> {
+        let mut checks = check_path(
+            Path::new("./resources/test/fixtures/D.py"),
+            &settings::Settings::for_rule(CheckCode::D204),
+            &fixer::Mode::Generate,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
     #[test]
     fn d205() -> Result<()> {
         let mut checks = check_path(
@@ -1056,6 +1104,18 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn d211() -> Result<()> {
+        let mut checks = check_path(
+            Path::new("./resources/test/fixtures/D.py"),
+            &settings::Settings::for_rule(CheckCode::D211),
+            &fixer::Mode::Generate,
+        )?;
+        checks.sort_by_key(|check| check.location);
+        insta::assert_yaml_snapshot!(checks);
+        Ok(())
+    }
+
     #[test]
     fn d212() -> Result<()> {
         let mut checks = check_path(

src/snapshots/ruff__linter__tests__d201.snap

@@ -0,0 +1,23 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    NoBlankLineBeforeFunction: 1
+  location:
+    row: 132
+    column: 5
+  end_location:
+    row: 132
+    column: 25
+  fix: ~
+- kind:
+    NoBlankLineBeforeFunction: 1
+  location:
+    row: 146
+    column: 5
+  end_location:
+    row: 146
+    column: 38
+  fix: ~
+

src/snapshots/ruff__linter__tests__d202.snap

@@ -0,0 +1,23 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    NoBlankLineAfterFunction: 1
+  location:
+    row: 137
+    column: 5
+  end_location:
+    row: 137
+    column: 25
+  fix: ~
+- kind:
+    NoBlankLineAfterFunction: 1
+  location:
+    row: 146
+    column: 5
+  end_location:
+    row: 146
+    column: 38
+  fix: ~
+

src/snapshots/ruff__linter__tests__d203.snap

@@ -0,0 +1,32 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    OneBlankLineBeforeClass: 0
+  location:
+    row: 156
+    column: 5
+  end_location:
+    row: 156
+    column: 33
+  fix: ~
+- kind:
+    OneBlankLineBeforeClass: 0
+  location:
+    row: 187
+    column: 5
+  end_location:
+    row: 187
+    column: 46
+  fix: ~
+- kind:
+    OneBlankLineBeforeClass: 0
+  location:
+    row: 521
+    column: 5
+  end_location:
+    row: 527
+    column: 8
+  fix: ~
+

src/snapshots/ruff__linter__tests__d204.snap

@@ -0,0 +1,23 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    OneBlankLineAfterClass: 0
+  location:
+    row: 176
+    column: 5
+  end_location:
+    row: 176
+    column: 25
+  fix: ~
+- kind:
+    OneBlankLineAfterClass: 0
+  location:
+    row: 187
+    column: 5
+  end_location:
+    row: 187
+    column: 46
+  fix: ~
+

src/snapshots/ruff__linter__tests__d211.snap

@@ -0,0 +1,23 @@
+---
+source: src/linter.rs
+expression: checks
+---
+- kind:
+    NoBlankLineBeforeClass: 1
+  location:
+    row: 165
+    column: 5
+  end_location:
+    row: 165
+    column: 30
+  fix: ~
+- kind:
+    NoBlankLineBeforeClass: 1
+  location:
+    row: 176
+    column: 5
+  end_location:
+    row: 176
+    column: 25
+  fix: ~
+