Fix more edge cases to do with generators

crates/ruff_linter/resources/test/fixtures/pydoclint/DOC201_numpy.py

@@ -193,3 +193,44 @@ def foo(s: str) -> str | None:
 def bar() -> int | None:
     """Bar-y method"""
     return
+
+
+from collections.abc import Iterator, Generator
+
+
+# This is okay -- returning `None` is implied by `Iterator[str]`;
+# no need to document it
+def generator_function() -> Iterator[str]:
+    """Generate some strings"""
+    yield from "abc"
+    return
+
+
+# This is okay -- returning `None` is stated by `Generator[str, None, None]`;
+# no need to document it
+def generator_function_2() -> Generator[str, None, None]:
+    """Generate some strings"""
+    yield from "abc"
+    return
+
+
+# DOC201 -- returns None but `Generator[str, None, int | None]`
+# indicates it could sometimes return `int`
+def generator_function_3() -> Generator[str, None, int | None]:
+    """Generate some strings"""
+    yield from "abc"
+    return
+
+
+# DOC201 -- no type annotation and a non-None return
+# indicates it could sometimes return `int`
+def generator_function_4():
+    """Generate some strings"""
+    yield from "abc"
+    return 42
+
+
+# DOC201 -- no `yield` expressions, so not a generator function
+def not_a_generator() -> Iterator[int]:
+    """"No returns documented here, oh no"""
+    return (x for x in range(42))

crates/ruff_linter/resources/test/fixtures/pydoclint/DOC402_google.py

@@ -89,11 +89,19 @@ def f(num: int):
     yield 1
 
 
-from collections import abc
+import collections.abc
 
 
 # DOC402
-def foo() -> abc.Generator[int | None, None, None]:
+def foo() -> collections.abc.Generator[int | None, None, None]:
+    """
+    Do something
+    """
+    yield
+
+
+# DOC402
+def bar() -> collections.abc.Iterator[int | None]:
     """
     Do something
     """

crates/ruff_linter/resources/test/fixtures/pydoclint/DOC402_numpy.py

@@ -130,3 +130,11 @@ def foo():
     """
     yield 1
     yield
+
+
+# DOC402
+def bar() -> typing.Iterator[int | None]:
+    """
+    Do something
+    """
+    yield

crates/ruff_linter/src/rules/pydoclint/rules/check_docstring.rs

@@ -732,26 +732,72 @@ fn yields_documented(
         || (matches!(convention, Some(Convention::Google)) && starts_with_yields(docstring))
 }
 
-/// Returns the first subscript element of a generator annotation, if it exists.
-fn generator_yield_annotation<'a>(expr: &'a Expr, checker: &'a Checker) -> Option<&'a Expr> {
-    let slice = expr.as_subscript_expr()?.slice.as_tuple_expr()?;
-    if matches!(
-        checker
-            .semantic()
-            .resolve_qualified_name(map_subscript(expr))?
-            .segments(),
-        [
-            "typing" | "typing_extensions",
-            "Generator" | "AsyncGenerator" | "Iterator" | "AsyncIterator"
-        ] | [
-            "collections",
-            "abc",
-            "Generator" | "AsyncGenerator" | "Iterator" | "AsyncIterator"
-        ]
-    ) {
-        return slice.elts.first();
+#[derive(Debug, Copy, Clone)]
+enum GeneratorOrIteratorArguments<'a> {
+    Unparameterized,
+    Single(&'a Expr),
+    Several(&'a [Expr]),
+}
+
+impl<'a> GeneratorOrIteratorArguments<'a> {
+    fn first(self) -> Option<&'a Expr> {
+        match self {
+            Self::Unparameterized => None,
+            Self::Single(element) => Some(element),
+            Self::Several(elements) => elements.first(),
+        }
+    }
+
+    fn indicates_none_returned(self) -> bool {
+        match self {
+            Self::Unparameterized => true,
+            Self::Single(_) => true,
+            Self::Several(elements) => elements.get(2).map_or(true, Expr::is_none_literal_expr),
+        }
+    }
+}
+
+/// Returns the arguments to a generator annotation, if it exists.
+fn generator_annotation_arguments<'a>(
+    expr: &'a Expr,
+    semantic: &'a SemanticModel,
+) -> Option<GeneratorOrIteratorArguments<'a>> {
+    let qualified_name = semantic.resolve_qualified_name(map_subscript(expr))?;
+    match qualified_name.segments() {
+        ["typing" | "typing_extensions", "Iterable" | "Iterator" | "AsyncIterator"]
+        | ["collections", "abc", "Iterable" | "Iterator" | "AsyncIterator"] => match expr {
+            Expr::Subscript(ast::ExprSubscript { slice, .. }) => {
+                Some(GeneratorOrIteratorArguments::Single(slice))
+            }
+            _ => Some(GeneratorOrIteratorArguments::Unparameterized),
+        },
+        ["typing" | "typing_extensions", "Generator" | "AsyncGenerator"]
+        | ["collections", "abc", "Generator" | "AsyncGenerator"] => match expr {
+            Expr::Subscript(ast::ExprSubscript { slice, .. }) => {
+                if let Expr::Tuple(tuple) = &**slice {
+                    Some(GeneratorOrIteratorArguments::Several(tuple.elts.as_slice()))
+                } else {
+                    // `Generator[int]` implies `Generator[int, None, None]`
+                    // as it uses a PEP-696 TypeVar with default values
+                    Some(GeneratorOrIteratorArguments::Single(slice))
+                }
+            }
+            _ => Some(GeneratorOrIteratorArguments::Unparameterized),
+        },
+        _ => None,
     }
-    None
+}
+
+fn is_generator_function_annotated_as_returning_none(
+    entries: &BodyEntries,
+    return_annotations: &Expr,
+    semantic: &SemanticModel,
+) -> bool {
+    if entries.yields.is_empty() {
+        return false;
+    }
+    generator_annotation_arguments(return_annotations, semantic)
+        .is_some_and(GeneratorOrIteratorArguments::indicates_none_returned)
 }
 
 /// DOC201, DOC202, DOC402, DOC403, DOC501, DOC502
@@ -769,8 +815,10 @@ pub(crate) fn check_docstring(
         return;
     };
 
+    let semantic = checker.semantic();
+
     // Ignore stubs.
-    if function_type::is_stub(function_def, checker.semantic()) {
+    if function_type::is_stub(function_def, semantic) {
         return;
     }
 
@@ -786,7 +834,7 @@ pub(crate) fn check_docstring(
     };
 
     let body_entries = {
-        let mut visitor = BodyVisitor::new(checker.semantic());
+        let mut visitor = BodyVisitor::new(semantic);
         visitor.visit_body(&function_def.body);
         visitor.finish()
     };
@@ -795,12 +843,26 @@ pub(crate) fn check_docstring(
     if checker.enabled(Rule::DocstringMissingReturns) {
         if !returns_documented(docstring, &docstring_sections, convention) {
             let extra_property_decorators = checker.settings.pydocstyle.property_decorators();
-            if !definition.is_property(extra_property_decorators, checker.semantic()) {
+            if !definition.is_property(extra_property_decorators, semantic) {
                 if let Some(body_return) = body_entries.returns.first() {
                     match function_def.returns.as_deref() {
-                        Some(returns) if !Expr::is_none_literal_expr(returns) => diagnostics.push(
-                            Diagnostic::new(DocstringMissingReturns, body_return.range()),
-                        ),
+                        Some(returns) => {
+                            // Ignore it if it's annotated as returning `None`
+                            // or it's a generator function annotated as returning `None`,
+                            // i.e. any of `-> None`, `-> Iterator[...]` or `-> Generator[..., ..., None]`
+                            if !returns.is_none_literal_expr()
+                                && !is_generator_function_annotated_as_returning_none(
+                                    &body_entries,
+                                    returns,
+                                    semantic,
+                                )
+                            {
+                                diagnostics.push(Diagnostic::new(
+                                    DocstringMissingReturns,
+                                    body_return.range(),
+                                ));
+                            }
+                        }
                         None if body_entries
                             .returns
                             .iter()
@@ -824,8 +886,9 @@ pub(crate) fn check_docstring(
             if let Some(body_yield) = body_entries.yields.first() {
                 match function_def.returns.as_deref() {
                     Some(returns)
-                        if generator_yield_annotation(returns, checker)
-                            .map_or(true, |expr| !Expr::is_none_literal_expr(expr)) =>
+                        if !generator_annotation_arguments(returns, semantic).is_some_and(
+                            |arguments| arguments.first().map_or(true, Expr::is_none_literal_expr),
+                        ) =>
                     {
                         diagnostics
                             .push(Diagnostic::new(DocstringMissingYields, body_yield.range()));
@@ -872,7 +935,7 @@ pub(crate) fn check_docstring(
 
     // Avoid applying "extraneous" rules to abstract methods. An abstract method's docstring _could_
     // document that it raises an exception without including the exception in the implementation.
-    if !visibility::is_abstract(&function_def.decorator_list, checker.semantic()) {
+    if !visibility::is_abstract(&function_def.decorator_list, semantic) {
         // DOC202
         if checker.enabled(Rule::DocstringExtraneousReturns) {
             if let Some(ref docstring_returns) = docstring_sections.returns {

crates/ruff_linter/src/rules/pydoclint/snapshots/ruff_linter__rules__pydoclint__tests__docstring-missing-returns_DOC201_numpy.py.snap

@@ -67,3 +67,30 @@ DOC201_numpy.py:195:5: DOC201 `return` is not documented in docstring
     |     ^^^^^^ DOC201
     |
     = help: Add a "Returns" section to the docstring
+
+DOC201_numpy.py:222:5: DOC201 `return` is not documented in docstring
+    |
+220 |     """Generate some strings"""
+221 |     yield from "abc"
+222 |     return
+    |     ^^^^^^ DOC201
+    |
+    = help: Add a "Returns" section to the docstring
+
+DOC201_numpy.py:230:5: DOC201 `return` is not documented in docstring
+    |
+228 |     """Generate some strings"""
+229 |     yield from "abc"
+230 |     return 42
+    |     ^^^^^^^^^ DOC201
+    |
+    = help: Add a "Returns" section to the docstring
+
+DOC201_numpy.py:236:5: DOC201 `return` is not documented in docstring
+    |
+234 | def not_a_generator() -> Iterator[int]:
+235 |     """"No returns documented here, oh no"""
+236 |     return (x for x in range(42))
+    |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ DOC201
+    |
+    = help: Add a "Returns" section to the docstring

crates/ruff_linter/src/rules/pydoclint/snapshots/ruff_linter__rules__pydoclint__tests__docstring-missing-yields_DOC402_google.py.snap

@@ -47,3 +47,12 @@ DOC402_google.py:100:5: DOC402 `yield` is not documented in docstring
     |     ^^^^^ DOC402
     |
     = help: Add a "Yields" section to the docstring
+
+DOC402_google.py:108:5: DOC402 `yield` is not documented in docstring
+    |
+106 |     Do something
+107 |     """
+108 |     yield
+    |     ^^^^^ DOC402
+    |
+    = help: Add a "Yields" section to the docstring

crates/ruff_linter/src/rules/pydoclint/snapshots/ruff_linter__rules__pydoclint__tests__docstring-missing-yields_DOC402_numpy.py.snap

@@ -57,3 +57,12 @@ DOC402_numpy.py:131:5: DOC402 `yield` is not documented in docstring
 132 |     yield
     |
     = help: Add a "Yields" section to the docstring
+
+DOC402_numpy.py:140:5: DOC402 `yield` is not documented in docstring
+    |
+138 |     Do something
+139 |     """
+140 |     yield
+    |     ^^^^^ DOC402
+    |
+    = help: Add a "Yields" section to the docstring