feat: Parse annotations in Iterator/Generator for Google docstrings

Issue #28: #28

docs/docstrings.md

@@ -47,20 +47,18 @@ Yields           | ✅     | ✅    | [❌][issue-section-sphinx-yields]
 
 ## Getting annotations/defaults from parent
 
-Section          | Google                             | Numpy                               | Sphinx
----------------- | ---------------------------------- | ----------------------------------- | ------
-Attributes       | ✅                                 | [❌][issue-parent-numpy-attributes] | [❌][issue-parent-sphinx-attributes]
-Deprecated       | /                                  | /                                   | /
-Examples         | /                                  | /                                   | /
-Other Parameters | ✅                                 | ✅                                  | [❌][issue-parent-sphinx-other-parameters]
-Parameters       | ✅                                 | ✅                                  | ✅
-Raises           | /                                  | /                                   | /
-Receives         | [❌][issue-parent-google-receives] | [❌][issue-parent-numpy-receives]   | [❌][issue-parent-sphinx-receives]
-Returns          | ✅                                 | [❌][issue-parent-numpy-returns]    | ✅
-Warns            | /                                  | /                                   | /
-Yields           | ✅                                 | [❌][issue-parent-numpy-yields]     | [❌][issue-parent-sphinx-yields]
-
-[issue-parent-google-receives]: https://github.com/mkdocstrings/griffe/issues/28
+Section          | Google | Numpy                               | Sphinx
+---------------- | ------ | ----------------------------------- | ------
+Attributes       | ✅     | [❌][issue-parent-numpy-attributes] | [❌][issue-parent-sphinx-attributes]
+Deprecated       | /      | /                                   | /
+Examples         | /      | /                                   | /
+Other Parameters | ✅     | ✅                                  | [❌][issue-parent-sphinx-other-parameters]
+Parameters       | ✅     | ✅                                  | ✅
+Raises           | /      | /                                   | /
+Receives         | ✅     | [❌][issue-parent-numpy-receives]   | [❌][issue-parent-sphinx-receives]
+Returns          | ✅     | [❌][issue-parent-numpy-returns]    | ✅
+Warns            | /      | /                                   | /
+Yields           | ✅     | [❌][issue-parent-numpy-yields]     | [❌][issue-parent-sphinx-yields]
 
 [issue-parent-numpy-attributes]: https://github.com/mkdocstrings/griffe/issues/29
 [issue-parent-numpy-receives]: https://github.com/mkdocstrings/griffe/issues/30

src/griffe/docstrings/google.py

@@ -347,11 +347,24 @@ def _read_returns_section(  # noqa: WPS231
             annotation = parse_annotation(annotation, docstring)
         else:
             # try to retrieve the annotation from the docstring parent
-            with suppress(AttributeError, KeyError):
+            with suppress(AttributeError, KeyError, ValueError):
                 annotation = docstring.parent.returns  # type: ignore[union-attr]
                 if len(block) > 1:
                     if annotation.is_tuple:
                         annotation = annotation.tuple_item(index)
+                    else:
+                        if annotation.is_iterator:
+                            return_item = annotation.iterator_item()
+                        elif annotation.is_generator:
+                            _, _, return_item = annotation.generator_items()
+                        else:
+                            raise ValueError
+                        if isinstance(return_item, Name):
+                            annotation = return_item
+                        elif return_item.is_tuple:
+                            annotation = return_item.tuple_item(index)
+                        else:
+                            annotation = return_item
 
             if annotation is None:
                 returned_value = repr(name) or index
@@ -388,12 +401,21 @@ def _read_yields_section(  # noqa: WPS231
             annotation = parse_annotation(annotation, docstring)
         else:
             # try to retrieve the annotation from the docstring parent
-            with suppress(AttributeError, KeyError):
+            with suppress(AttributeError, KeyError, ValueError):
                 annotation = docstring.parent.returns  # type: ignore[union-attr]
-                # TODO: support getting yield part and exploding tuple (in a generator/iterator)
-                # if len(block) > 1:
-                #     if annotation.is_tuple:
-                #         annotation = annotation.tuple_item(index)
+                if len(block) > 1:
+                    if annotation.is_iterator:
+                        yield_item = annotation.iterator_item()
+                    elif annotation.is_generator:
+                        yield_item, _, _ = annotation.generator_items()
+                    else:
+                        raise ValueError
+                    if isinstance(yield_item, Name):
+                        annotation = yield_item
+                    elif yield_item.is_tuple:
+                        annotation = yield_item.tuple_item(index)
+                    else:
+                        annotation = yield_item
 
             if annotation is None:
                 yielded_value = repr(name) or index
@@ -428,14 +450,18 @@ def _read_receives_section(  # noqa: WPS231
         if annotation:
             # try to compile the annotation to transform it into an expression
             annotation = parse_annotation(annotation, docstring)
-        # else:
-        # try to retrieve the annotation from the docstring parent
-        # TODO: support getting receive part and exploding tuple (in a generator/iterator)
-        # with suppress(AttributeError, KeyError):
-        #     annotation = docstring.parent.returns  # type: ignore[union-attr]
-        #     if len(block) > 1:
-        #         if annotation.is_tuple:
-        #             annotation = annotation.tuple_item(index)
+        else:
+            # try to retrieve the annotation from the docstring parent
+            with suppress(AttributeError, KeyError):
+                annotation = docstring.parent.returns  # type: ignore[union-attr]
+                if len(block) > 1 and annotation.is_generator:
+                    _, receives_item, _ = annotation.generator_items()
+                    if isinstance(receives_item, Name):
+                        annotation = receives_item
+                    elif receives_item.is_tuple:
+                        annotation = receives_item.tuple_item(index)
+                    else:
+                        annotation = receives_item
 
         if annotation is None:
             received_value = repr(name) or index

src/griffe/expressions.py

@@ -118,14 +118,41 @@ def full(self) -> str:
         """
         return str(self)
 
+    @property
+    def kind(self) -> str:
+        """Return the main type object as a string.
+
+        Returns:
+            The main type of this expression.
+        """
+        return str(self).split("[", 1)[0].rsplit(".", 1)[-1].lower()
+
     @property
     def is_tuple(self) -> bool:
         """Tell whether this expression represents a tuple.
 
         Returns:
             True or False.
         """
-        return str(self).split("[", 1)[0].rsplit(".", 1)[-1].lower() == "tuple"
+        return self.kind == "tuple"
+
+    @property
+    def is_iterator(self) -> bool:
+        """Tell whether this expression represents an iterator.
+
+        Returns:
+            True or False.
+        """
+        return self.kind == "iterator"
+
+    @property
+    def is_generator(self) -> bool:
+        """Tell whether this expression represents a generator.
+
+        Returns:
+            True or False.
+        """
+        return self.kind == "generator"
 
     def tuple_item(self, nth: int) -> str | Name:
         """Return the n-th item of this tuple expression.
@@ -141,3 +168,31 @@ def tuple_item(self, nth: int) -> str | Name:
         #       N , N , N
         #       0 1 2 3 4
         return self[2][2 * nth]
+
+    def tuple_items(self) -> list[Name | Expression]:
+        """Return a tuple items as a list.
+
+        Returns:
+            The tuple items.
+        """
+        return self[2][::2]
+
+    def iterator_item(self) -> Name | Expression:
+        """Return the item of an iterator.
+
+        Returns:
+            The iterator item.
+        """
+        # Iterator[ItemType]
+        return self[2]
+
+    def generator_items(self) -> tuple[Name | Expression, Name | Expression, Name | Expression]:
+        """Return the items of a generator.
+
+        Returns:
+            The yield type.
+            The send/receive type.
+            The return type.
+        """
+        # Generator[Yield, Send/Receive, Return]
+        return self[2][0], self[2][2], self[2][4]

tests/test_docstrings/test_google.py

@@ -6,8 +6,9 @@
 
 import pytest
 
-from griffe.dataclasses import Class, Function, Module, Parameter, Parameters
+from griffe.dataclasses import Class, Docstring, Function, Module, Parameter, Parameters
 from griffe.docstrings.dataclasses import DocstringSectionKind
+from griffe.docstrings.utils import parse_annotation
 
 
 # =============================================================================================
@@ -757,6 +758,99 @@ def test_parse_yields_section_with_return_annotation(parse_google):
     assert not warnings
 
 
+@pytest.mark.parametrize(
+    "return_annotation",
+    [
+        "Iterator[tuple[int, float]]",
+        "Generator[tuple[int, float], ..., ...]",
+    ],
+)
+def test_parse_yields_tuple_in_iterator_or_generator(parse_google, return_annotation):
+    """Parse Yields annotations in Iterator or Generator types.
+
+    Parameters:
+        parse_google: Fixture parser.
+        return_annotation: Parametrized return annotation as a string.
+    """
+    docstring = """
+        Summary.
+
+        Yields:
+            a: Whatever.
+            b: Whatever.
+    """
+    sections, _ = parse_google(
+        docstring,
+        parent=Function(
+            "func",
+            returns=parse_annotation(return_annotation, Docstring("d", parent=Function("f"))),
+        ),
+    )
+    yields = sections[1].value
+    assert yields[0].name == "a"
+    assert yields[0].annotation.source == "int"
+    assert yields[1].name == "b"
+    assert yields[1].annotation.source == "float"
+
+
+# =============================================================================================
+# Receives sections
+def test_parse_receives_tuple_in_generator(parse_google):
+    """Parse Receives annotations in Generator type.
+
+    Parameters:
+        parse_google: Fixture parser.
+    """
+    docstring = """
+        Summary.
+
+        Receives:
+            a: Whatever.
+            b: Whatever.
+    """
+    sections, _ = parse_google(
+        docstring,
+        parent=Function(
+            "func",
+            returns=parse_annotation("Generator[..., tuple[int, float], ...]", Docstring("d", parent=Function("f"))),
+        ),
+    )
+    receives = sections[1].value
+    assert receives[0].name == "a"
+    assert receives[0].annotation.source == "int"
+    assert receives[1].name == "b"
+    assert receives[1].annotation.source == "float"
+
+
+# =============================================================================================
+# Returns sections
+def test_parse_returns_tuple_in_generator(parse_google):
+    """Parse Returns annotations in Generator type.
+
+    Parameters:
+        parse_google: Fixture parser.
+    """
+    docstring = """
+        Summary.
+
+        Returns:
+            a: Whatever.
+            b: Whatever.
+    """
+    sections, _ = parse_google(
+        docstring,
+        parent=Function(
+            "func",
+            returns=parse_annotation("Generator[..., ..., tuple[int, float]]", Docstring("d", parent=Function("f"))),
+        ),
+    )
+    returns = sections[1].value
+    assert returns[0].name == "a"
+    assert returns[0].annotation.source == "int"
+    assert returns[1].name == "b"
+    assert returns[1].annotation.source == "float"
+
+
 # =============================================================================================
 # Parser special features
 def test_parse_admonitions(parse_google):