Enable definition tracking for docstrings (#407)

src/check_ast.rs

@@ -21,7 +21,7 @@ use crate::ast::visitor::{walk_excepthandler, Visitor};
 use crate::ast::{checkers, helpers, operations, visitor};
 use crate::autofix::{fixer, fixes};
 use crate::checks::{Check, CheckCode, CheckKind};
-use crate::docstrings::{Docstring, DocstringKind};
+use crate::docstrings::{Definition, DefinitionKind, Documentable};
 use crate::python::builtins::{BUILTINS, MAGIC_GLOBALS};
 use crate::python::future::ALL_FEATURE_NAMES;
 use crate::settings::{PythonVersion, Settings};
@@ -31,7 +31,7 @@ pub const GLOBAL_SCOPE_INDEX: usize = 0;
 
 pub struct Checker<'a> {
     // Input data.
-    path: &'a Path,
+    pub(crate) path: &'a Path,
     // TODO(charlie): Separate immutable from mutable state. (None of these should ever change.)
     pub(crate) locator: SourceCodeLocator<'a>,
     pub(crate) settings: &'a Settings,
@@ -39,7 +39,7 @@ pub struct Checker<'a> {
     // Computed checks.
     checks: Vec<Check>,
     // Docstring tracking.
-    docstrings: Vec<Docstring<'a>>,
+    docstrings: Vec<Definition<'a>>,
     // Edit tracking.
     // TODO(charlie): Instead of exposing deletions, wrap in a public API.
     pub(crate) deletions: BTreeSet<usize>,
@@ -125,36 +125,8 @@ where
             StmtKind::Import { .. } => {
                 self.futures_allowed = false;
             }
-            StmtKind::Expr { value } => {
-                // Track all docstrings: module-, class-, and function-level.
-                let mut is_module_docstring = false;
-                if matches!(
-                    &value.node,
-                    ExprKind::Constant {
-                        value: Constant::Str(_),
-                        ..
-                    }
-                ) {
-                    if let Some(docstring) = docstrings::extract(self, stmt, value) {
-                        if matches!(&docstring.kind, DocstringKind::Module) {
-                            is_module_docstring = true;
-                        }
-                        self.docstrings.push(docstring);
-                    }
-                }
-
-                if !is_module_docstring {
-                    if !self.seen_import_boundary
-                        && !operations::in_nested_block(&self.parent_stack, &self.parents)
-                    {
-                        self.seen_import_boundary = true;
-                    }
-                    self.futures_allowed = false;
-                }
-            }
             node => {
                 self.futures_allowed = false;
-
                 if !self.seen_import_boundary
                     && !helpers::is_assignment_to_a_dunder(node)
                     && !operations::in_nested_block(&self.parent_stack, &self.parents)
@@ -594,7 +566,23 @@ where
 
         // Recurse.
         match &stmt.node {
-            StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
+            StmtKind::FunctionDef { body, .. } | StmtKind::AsyncFunctionDef { body, .. } => {
+                // TODO(charlie): Track public / private.
+                // Grab all parents, to enable nested definition tracking. (Ignore the most recent
+                // parent which, confusingly, is `stmt`.)
+                let parents = self
+                    .parent_stack
+                    .iter()
+                    .take(self.parents.len() - 1)
+                    .map(|index| self.parents[*index])
+                    .collect();
+                self.docstrings.push(docstrings::extract(
+                    parents,
+                    stmt,
+                    body,
+                    Documentable::Function,
+                ));
+
                 self.deferred_functions.push((
                     stmt,
                     self.scope_stack.clone(),
@@ -602,6 +590,22 @@ where
                 ));
             }
             StmtKind::ClassDef { body, .. } => {
+                // TODO(charlie): Track public / priva
+                // Grab all parents, to enable nested definition tracking. (Ignore the most recent
+                // parent which, confusingly, is `stmt`.)
+                let parents = self
+                    .parent_stack
+                    .iter()
+                    .take(self.parents.len() - 1)
+                    .map(|index| self.parents[*index])
+                    .collect();
+                self.docstrings.push(docstrings::extract(
+                    parents,
+                    stmt,
+                    body,
+                    Documentable::Class,
+                ));
+
                 for stmt in body {
                     self.visit_stmt(stmt);
                 }
@@ -1643,6 +1647,22 @@ impl<'a> Checker<'a> {
         }
     }
 
+    fn visit_docstring<'b>(&mut self, python_ast: &'b Suite) -> bool
+    where
+        'b: 'a,
+    {
+        let docstring = docstrings::docstring_from(python_ast);
+        self.docstrings.push(Definition {
+            kind: if self.path.ends_with("__init__.py") {
+                DefinitionKind::Package
+            } else {
+                DefinitionKind::Module
+            },
+            docstring,
+        });
+        docstring.is_some()
+    }
+
     fn check_deferred_annotations(&mut self) {
         while let Some((expr, scopes, parents)) = self.deferred_annotations.pop() {
             self.parent_stack = parents;
@@ -1988,6 +2008,13 @@ pub fn check_ast(
     checker.push_scope(Scope::new(ScopeKind::Module));
     checker.bind_builtins();
 
+    // Check for module docstring.
+    let python_ast = if checker.visit_docstring(python_ast) {
+        &python_ast[1..]
+    } else {
+        python_ast
+    };
+
     // Iterate over the AST.
     for stmt in python_ast {
         checker.visit_stmt(stmt);

src/docstrings.rs

@@ -7,77 +7,101 @@ use crate::check_ast::Checker;
 use crate::checks::{Check, CheckCode, CheckKind};
 
 #[derive(Debug)]
-pub enum DocstringKind<'a> {
+pub enum DefinitionKind<'a> {
     Module,
-    Function(&'a Stmt),
+    Package,
     Class(&'a Stmt),
+    NestedClass(&'a Stmt),
+    Function(&'a Stmt),
+    NestedFunction(&'a Stmt),
+    Method(&'a Stmt),
 }
 
 #[derive(Debug)]
-pub struct Docstring<'a> {
+pub struct Definition<'a> {
-    pub kind: DocstringKind<'a>,
+    pub kind: DefinitionKind<'a>,
-    pub expr: &'a Expr,
+    pub docstring: Option<&'a Expr>,
 }
 
-/// Extract a `Docstring` from an `Expr`.
+pub enum Documentable {
-pub fn extract<'a, 'b>(
+    Class,
-    checker: &'a Checker<'b>,
+    Function,
-    stmt: &'b Stmt,
-    expr: &'b Expr,
-) -> Option<Docstring<'b>> {
-    let defined_in = checker
-        .binding_context()
-        .defined_in
-        .map(|index| checker.parents[index]);
-    match defined_in {
-        None => {
-            if checker.initial {
-                return Some(Docstring {
-                    kind: DocstringKind::Module,
-                    expr,
-                });
-            }
-        }
-        Some(parent) => {
-            if let StmtKind::FunctionDef { body, .. }
-            | StmtKind::AsyncFunctionDef { body, .. }
-            | StmtKind::ClassDef { body, .. } = &parent.node
-            {
-                if body.first().map(|node| node == stmt).unwrap_or_default() {
-                    return Some(Docstring {
-                        kind: if matches!(&parent.node, StmtKind::ClassDef { .. }) {
-                            DocstringKind::Class(parent)
-                        } else {
-                            DocstringKind::Function(parent)
-                        },
-                        expr,
-                    });
-                }
-            }
-        }
 }
 
+fn nest(parents: &[&Stmt]) -> Option<Documentable> {
+    for parent in parents.iter().rev() {
+        match &parent.node {
+            StmtKind::FunctionDef { .. } | StmtKind::AsyncFunctionDef { .. } => {
+                return Some(Documentable::Function)
+            }
+            StmtKind::ClassDef { .. } => return Some(Documentable::Class),
+            _ => {}
+        }
+    }
     None
 }
 
-/// Extract the source code range for a `Docstring`.
+/// Extract a docstring from a function or class body.
-fn range_for(docstring: &Docstring) -> Range {
+pub fn docstring_from(suite: &[Stmt]) -> Option<&Expr> {
+    if let Some(stmt) = suite.first() {
+        if let StmtKind::Expr { value } = &stmt.node {
+            if matches!(
+                &value.node,
+                ExprKind::Constant {
+                    value: Constant::Str(_),
+                    ..
+                }
+            ) {
+                return Some(value);
+            }
+        }
+    }
+    None
+}
+
+/// Extract a `Definition` from the AST node defined by a `Stmt`.
+pub fn extract<'a>(
+    parents: Vec<&'a Stmt>,
+    stmt: &'a Stmt,
+    body: &'a [Stmt],
+    kind: Documentable,
+) -> Definition<'a> {
+    let expr = docstring_from(body);
+    match kind {
+        Documentable::Function => Definition {
+            kind: match nest(&parents) {
+                None => DefinitionKind::Function(stmt),
+                Some(Documentable::Function) => DefinitionKind::NestedFunction(stmt),
+                Some(Documentable::Class) => DefinitionKind::Method(stmt),
+            },
+            docstring: expr,
+        },
+        Documentable::Class => Definition {
+            kind: match nest(&parents) {
+                None => DefinitionKind::Class(stmt),
+                Some(_) => DefinitionKind::NestedClass(stmt),
+            },
+            docstring: expr,
+        },
+    }
+}
+
+/// Extract the source code range for a docstring.
+fn range_for(docstring: &Expr) -> Range {
     // RustPython currently omits the first quotation mark in a string, so offset the location.
     Range {
-        location: Location::new(
+        location: Location::new(docstring.location.row(), docstring.location.column() - 1),
-            docstring.expr.location.row(),
+        end_location: docstring.end_location,
-            docstring.expr.location.column() - 1,
-        ),
-        end_location: docstring.expr.end_location,
     }
 }
 
 /// D200
-pub fn one_liner(checker: &mut Checker, docstring: &Docstring) {
+pub fn one_liner(checker: &mut Checker, definition: &Definition) {
+    if let Some(docstring) = &definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             let mut line_count = 0;
             let mut non_empty_line_count = 0;
@@ -96,6 +120,7 @@ pub fn one_liner(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 static COMMENT_REGEX: Lazy<Regex> = Lazy::new(|| Regex::new(r"^\s*#").unwrap());
 
@@ -103,12 +128,16 @@ 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) {
+pub fn blank_before_after_function(checker: &mut Checker, definition: &Definition) {
-    if let DocstringKind::Function(parent) = &docstring.kind {
+    if let Some(docstring) = definition.docstring {
+        if let DefinitionKind::Function(parent)
+        | DefinitionKind::NestedFunction(parent)
+        | DefinitionKind::Method(parent) = &definition.kind
+        {
             if let ExprKind::Constant {
                 value: Constant::Str(_),
                 ..
-        } = &docstring.expr.node
+            } = &docstring.node
             {
                 let (before, _, after) = checker
                     .locator
@@ -144,7 +173,8 @@ pub fn blank_before_after_function(checker: &mut Checker, docstring: &Docstring)
                     // class.
                     if !all_blank_after
                         && blank_lines_after != 0
-                    && !(blank_lines_after == 1 && INNER_FUNCTION_OR_CLASS_REGEX.is_match(after))
+                        && !(blank_lines_after == 1
+                            && INNER_FUNCTION_OR_CLASS_REGEX.is_match(after))
                     {
                         checker.add_check(Check::new(
                             CheckKind::NoBlankLineAfterFunction(blank_lines_after),
@@ -155,14 +185,18 @@ pub fn blank_before_after_function(checker: &mut Checker, docstring: &Docstring)
             }
         }
     }
+}
 
 /// D203, D204, D211
-pub fn blank_before_after_class(checker: &mut Checker, docstring: &Docstring) {
+pub fn blank_before_after_class(checker: &mut Checker, definition: &Definition) {
-    if let DocstringKind::Class(parent) = &docstring.kind {
+    if let Some(docstring) = &definition.docstring {
+        if let DefinitionKind::Class(parent) | DefinitionKind::NestedClass(parent) =
+            &definition.kind
+        {
             if let ExprKind::Constant {
                 value: Constant::Str(_),
                 ..
-        } = &docstring.expr.node
+            } = &docstring.node
             {
                 let (before, _, after) = checker
                     .locator
@@ -177,13 +211,17 @@ pub fn blank_before_after_class(checker: &mut Checker, docstring: &Docstring) {
                         .skip(1)
                         .take_while(|line| line.trim().is_empty())
                         .count();
-                if blank_lines_before != 0 && checker.settings.enabled.contains(&CheckCode::D211) {
+                    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) {
+                    if blank_lines_before != 1
+                        && checker.settings.enabled.contains(&CheckCode::D203)
+                    {
                         checker.add_check(Check::new(
                             CheckKind::OneBlankLineBeforeClass(blank_lines_before),
                             range_for(docstring),
@@ -211,13 +249,15 @@ pub fn blank_before_after_class(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 /// D205
-pub fn blank_after_summary(checker: &mut Checker, docstring: &Docstring) {
+pub fn blank_after_summary(checker: &mut Checker, definition: &Definition) {
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             let mut lines_count = 1;
             let mut blanks_count = 0;
@@ -237,13 +277,15 @@ pub fn blank_after_summary(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 /// D209
-pub fn newline_after_last_paragraph(checker: &mut Checker, docstring: &Docstring) {
+pub fn newline_after_last_paragraph(checker: &mut Checker, definition: &Definition) {
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             let mut line_count = 0;
             for line in string.lines() {
@@ -268,13 +310,15 @@ pub fn newline_after_last_paragraph(checker: &mut Checker, docstring: &Docstring
             }
         }
     }
+}
 
 /// D210
-pub fn no_surrounding_whitespace(checker: &mut Checker, docstring: &Docstring) {
+pub fn no_surrounding_whitespace(checker: &mut Checker, definition: &Definition) {
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             let mut lines = string.lines();
             if let Some(line) = lines.next() {
@@ -290,13 +334,15 @@ pub fn no_surrounding_whitespace(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 /// D212, D213
-pub fn multi_line_summary_start(checker: &mut Checker, docstring: &Docstring) {
+pub fn multi_line_summary_start(checker: &mut Checker, definition: &Definition) {
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             if string.lines().nth(1).is_some() {
                 let content = checker
@@ -321,13 +367,15 @@ pub fn multi_line_summary_start(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 /// D300
-pub fn triple_quotes(checker: &mut Checker, docstring: &Docstring) {
+pub fn triple_quotes(checker: &mut Checker, definition: &Definition) {
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             let content = checker
                 .locator
@@ -347,13 +395,15 @@ pub fn triple_quotes(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 /// D400
-pub fn ends_with_period(checker: &mut Checker, docstring: &Docstring) {
+pub fn ends_with_period(checker: &mut Checker, definition: &Definition) {
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             if let Some(string) = string.lines().next() {
                 if !string.ends_with('.') {
@@ -362,19 +412,28 @@ pub fn ends_with_period(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 /// D402
-pub fn no_signature(checker: &mut Checker, docstring: &Docstring) {
+pub fn no_signature(checker: &mut Checker, definition: &Definition) {
-    if let DocstringKind::Function(parent) = docstring.kind {
+    if let Some(docstring) = definition.docstring {
+        if let DefinitionKind::Function(parent)
+        | DefinitionKind::NestedFunction(parent)
+        | DefinitionKind::Method(parent) = definition.kind
+        {
             if let StmtKind::FunctionDef { name, .. } = &parent.node {
                 if let ExprKind::Constant {
                     value: Constant::Str(string),
                     ..
-            } = &docstring.expr.node
+                } = &docstring.node
                 {
                     if let Some(first_line) = string.lines().next() {
                         if first_line.contains(&format!("{name}(")) {
-                        checker.add_check(Check::new(CheckKind::NoSignature, range_for(docstring)));
+                            checker.add_check(Check::new(
+                                CheckKind::NoSignature,
+                                range_for(docstring),
+                            ));
+                        }
                     }
                 }
             }
@@ -383,15 +442,16 @@ pub fn no_signature(checker: &mut Checker, docstring: &Docstring) {
 }
 
 /// D403
-pub fn capitalized(checker: &mut Checker, docstring: &Docstring) {
+pub fn capitalized(checker: &mut Checker, definition: &Definition) {
-    if !matches!(docstring.kind, DocstringKind::Function(_)) {
+    if !matches!(definition.kind, DefinitionKind::Function(_)) {
         return;
     }
 
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             if let Some(first_word) = string.split(' ').next() {
                 if first_word == first_word.to_uppercase() {
@@ -413,13 +473,15 @@ pub fn capitalized(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 /// D415
-pub fn ends_with_punctuation(checker: &mut Checker, docstring: &Docstring) {
+pub fn ends_with_punctuation(checker: &mut Checker, definition: &Definition) {
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             if let Some(string) = string.lines().next() {
                 if !(string.ends_with('.') || string.ends_with('!') || string.ends_with('?')) {
@@ -431,13 +493,15 @@ pub fn ends_with_punctuation(checker: &mut Checker, docstring: &Docstring) {
             }
         }
     }
+}
 
 /// D419
-pub fn not_empty(checker: &mut Checker, docstring: &Docstring) -> bool {
+pub fn not_empty(checker: &mut Checker, definition: &Definition) -> bool {
+    if let Some(docstring) = definition.docstring {
         if let ExprKind::Constant {
             value: Constant::Str(string),
             ..
-    } = &docstring.expr.node
+        } = &docstring.node
         {
             if string.trim().is_empty() {
                 if checker.settings.enabled.contains(&CheckCode::D419) {
@@ -446,5 +510,6 @@ pub fn not_empty(checker: &mut Checker, docstring: &Docstring) -> bool {
                 return false;
             }
         }
+    }
     true
 }