Annotate tuple return types (#10)

Author: flying-sheep

scanpydoc/elegant_typehints/__init__.py

@@ -82,4 +82,9 @@ def setup(app: Sphinx) -> Dict[str, Any]:
             name, partial(_role_annot, additional_classes=name.split("-"))
         )
 
+    from .return_tuple import process_docstring  # , process_signature
+
+    app.connect("autodoc-process-docstring", process_docstring)
+    # app.connect("autodoc-process-signature", process_signature)
+
     return metadata

scanpydoc/elegant_typehints/formatting.py

@@ -112,10 +112,7 @@ def format_annotation(annotation: Type[Any], fully_qualified: bool = False) -> s
     curframe = inspect.currentframe()
     calframe = inspect.getouterframes(curframe, 2)
     if calframe[1][3] == "process_docstring":
-        annot_fmt = (
-            f":annotation-terse:`{_escape(_format_terse(annotation, fully_qualified))}`\\ "
-            f":annotation-full:`{_escape(_format_full(annotation, fully_qualified))}`"
-        )
+        annot_fmt = format_both(annotation, fully_qualified)
         if elegant_typehints.annotate_defaults:
             variables = calframe[1].frame.f_locals
             sig = inspect.signature(variables["obj"])
@@ -128,6 +125,13 @@ def format_annotation(annotation: Type[Any], fully_qualified: bool = False) -> s
         return _format_full(annotation, fully_qualified)
 
 
+def format_both(annotation: Type[Any], fully_qualified: bool = False):
+    return (
+        f":annotation-terse:`{_escape(_format_terse(annotation, fully_qualified))}`\\ "
+        f":annotation-full:`{_escape(_format_full(annotation, fully_qualified))}`"
+    )
+
+
 def _role_annot(
     name: str,
     rawtext: str,

scanpydoc/elegant_typehints/return_tuple.py

@@ -0,0 +1,89 @@
+import inspect
+import re
+from typing import get_type_hints, Any, Union, Optional, Type, Tuple, List
+
+from sphinx.application import Sphinx
+from sphinx.ext.autodoc import Options
+
+from .formatting import format_both
+
+
+re_ret = re.compile("^:returns?: ")
+
+
+def get_tuple_annot(annotation: Optional[Type]) -> Optional[Tuple[Type, ...]]:
+    if annotation is None:
+        return None
+    origin = getattr(annotation, "__origin__")
+    if not origin:
+        return None
+    if origin is Union:
+        for annotation in annotation.__args__:
+            origin = getattr(annotation, "__origin__")
+            if origin in (tuple, Tuple):
+                break
+        else:
+            return None
+    return annotation.__args__
+
+
+def process_docstring(
+    app: Sphinx,
+    what: str,
+    name: str,
+    obj: Any,
+    options: Optional[Options],
+    lines: List[str],
+) -> None:
+    # Handle complex objects
+    if isinstance(obj, property):
+        obj = obj.fget
+    if not callable(obj):
+        return
+    if what in ("class", "exception"):
+        obj = obj.__init__
+    obj = inspect.unwrap(obj)
+    ret_types = get_tuple_annot(get_type_hints(obj).get("return"))
+    if ret_types is None:
+        return
+
+    # Get return section
+    i_prefix = None
+    l_start = None
+    l_end = None
+    for l, line in enumerate(lines):
+        if i_prefix is None:
+            m = re_ret.match(line)
+            if m:
+                i_prefix = m.span()[1]
+                l_start = l
+        elif len(line[:i_prefix].strip()) > 0:
+            l_end = l - 1
+            break
+    else:
+        l_end = len(lines) - 1
+    if i_prefix is None:
+        return
+
+    # Meat
+    idxs_ret_names = []
+    for l, line in enumerate([l[i_prefix:] for l in lines[l_start : l_end + 1]]):
+        if line.isidentifier() and lines[l + l_start + 1].startswith("    "):
+            idxs_ret_names.append(l + l_start)
+
+    if len(idxs_ret_names) == len(ret_types):
+        for l, rt in zip(idxs_ret_names, ret_types):
+            typ = format_both(rt, app.config.typehints_fully_qualified)
+            lines[l : l + 1] = [f"{lines[l]} : {typ}"]
+
+
+# def process_signature(
+#     app: Sphinx,
+#     what: str,
+#     name: str,
+#     obj: Any,
+#     options: Options,
+#     signature: Optional[str],
+#     return_annotation: str,
+# ) -> Optional[Tuple[Optional[str], Optional[str]]]:
+#     return signature, return_annotation

tests/test_elegant_typehints.py

@@ -1,4 +1,5 @@
 import inspect
+import re
 import typing as t
 
 try:
@@ -7,14 +8,16 @@
     from typing_extensions import Literal
 
 import pytest
+import sphinx_autodoc_typehints as sat
 from sphinx.application import Sphinx
-from sphinx_autodoc_typehints import process_docstring
 
 from scanpydoc.elegant_typehints.formatting import (
     format_annotation,
     _format_terse,
     _format_full,
 )
+from scanpydoc.elegant_typehints.return_tuple import process_docstring
+
 
 TestCls = type("Class", (), {})
 TestCls.__module__ = "_testmod"
@@ -23,7 +26,11 @@
 @pytest.fixture
 def app(make_app_setup) -> Sphinx:
     return make_app_setup(
-        extensions="scanpydoc.elegant_typehints",
+        extensions=[
+            "sphinx.ext.napoleon",
+            "sphinx_autodoc_typehints",
+            "scanpydoc.elegant_typehints",
+        ],
         qualname_overrides={"_testmod.Class": "test.Class"},
     )
 
@@ -34,6 +41,7 @@ def process_doc(app):
 
     def process(fn: t.Callable) -> t.List[str]:
         lines = inspect.getdoc(fn).split("\n")
+        sat.process_docstring(app, "function", fn.__name__, fn, None, lines)
         process_docstring(app, "function", fn.__name__, fn, None, lines)
         return lines
 
@@ -202,3 +210,69 @@ def test_typing_class_nested(app):
         ":py:data:`~typing.Tuple`\\[:py:class:`int`, :py:class:`str`]"
         "]"
     )
+
+
+@pytest.mark.parametrize(
+    "docstring",
+    [
+        """
+        :param: x
+        :return: foo
+                     A foo!
+                 bar
+                     A bar!
+        """,
+        """
+        :return: foo
+                     A foo!
+                 bar
+                     A bar!
+        :param: x
+        """,
+    ],
+    ids=["Last", "First"],
+)
+@pytest.mark.parametrize(
+    "return_ann",
+    [t.Tuple[str, int], t.Optional[t.Tuple[str, int]]],
+    ids=["Tuple", "Optional[Tuple]"],
+)
+def test_return(process_doc, docstring, return_ann):
+    def fn_test():
+        pass
+
+    fn_test.__doc__ = docstring
+    fn_test.__annotations__["return"] = return_ann
+    lines = [
+        l
+        for l in process_doc(fn_test)
+        if not re.match("^:(rtype|param|annotation-(full|terse)):", l)
+    ]
+    assert lines == [
+        r":return: foo : "
+        r":annotation-terse:`:py:class:\`str\``\ "
+        r":annotation-full:`:py:class:\`str\``",
+        "             A foo!",
+        r"         bar : "
+        r":annotation-terse:`:py:class:\`int\``\ "
+        r":annotation-full:`:py:class:\`int\``",
+        "             A bar!",
+    ]
+
+
+def test_return_too_many(process_doc):
+    def fn_test() -> t.Tuple[int, str]:
+        """
+        :return: foo
+                     A foo!
+                 bar
+                     A bar!
+                 baz
+                     A baz!
+        """
+
+    assert not any(
+        "annotation-terse" in l
+        for l in process_doc(fn_test)
+        if not l.startswith(":rtype:")
+    )