Configurable formatter (#171)

Co-authored-by: Bernát Gábor <[email protected]>

Author: flying-sheep

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - Resolve type guard imports before evaluating annotations for objects
 - Fix crash when the `inspect` module returns an invalid python syntax source
+- Made formatting function configurable using the option `typehints_formatter`
 
 ## 1.14.1
 

README.md

@@ -72,6 +72,9 @@ The following configuration options are accepted:
   None\]). If `False`, the \"Optional\"-type is kept. Note: If `False`, **any** Union containing `None` will be
   displayed as Optional! Note: If an optional parameter has only a single type (e.g Optional\[A\] or Union\[A, None\]),
   it will **always** be displayed as Optional!
+- `typehints_formatter` (default: `None`): If set to a function, this function will be called with `annotation` as first
+  argument and `sphinx.config.Config` argument second. The function is expected to return a string with reStructuredText
+  code or `None` to fall back to the default formatter.
 
 ## How it works
 

src/sphinx_autodoc_typehints/__init__.py

@@ -6,10 +6,10 @@
 import textwrap
 import typing
 from ast import FunctionDef, Module, stmt
-from functools import partial
-from typing import Any, AnyStr, NewType, TypeVar, get_type_hints
+from typing import Any, AnyStr, Callable, NewType, TypeVar, get_type_hints
 
 from sphinx.application import Sphinx
+from sphinx.config import Config
 from sphinx.environment import BuildEnvironment
 from sphinx.ext.autodoc import Options
 from sphinx.util import logging
@@ -90,7 +90,13 @@ def get_annotation_args(annotation: Any, module: str, class_name: str) -> tuple[
     return getattr(annotation, "__args__", ())
 
 
-def format_annotation(annotation: Any, fully_qualified: bool = False, simplify_optional_unions: bool = True) -> str:
+def format_annotation(annotation: Any, config: Config) -> str:
+    typehints_formatter: Callable[..., str] | None = getattr(config, "typehints_formatter", None)
+    if typehints_formatter is not None:
+        formatted = typehints_formatter(annotation, config)
+        if formatted is not None:
+            return formatted
+
     # Special cases
     if annotation is None or annotation is type(None):  # noqa: E721
         return ":py:obj:`None`"
@@ -116,6 +122,7 @@ def format_annotation(annotation: Any, fully_qualified: bool = False, simplify_o
         module = "typing"
 
     full_name = f"{module}.{class_name}" if module != "builtins" else class_name
+    fully_qualified: bool = getattr(config, "fully_qualified", False)
     prefix = "" if fully_qualified or full_name == class_name else "~"
     role = "data" if class_name in _PYDATA_ANNOTATIONS else "class"
     args_format = "\\[{}]"
@@ -131,20 +138,21 @@ def format_annotation(annotation: Any, fully_qualified: bool = False, simplify_o
         if len(args) == 2:
             full_name = "typing.Optional"
             args = tuple(x for x in args if x is not type(None))  # noqa: E721
-        elif not simplify_optional_unions:
-            full_name = "typing.Optional"
-            args_format = f"\\[:py:data:`{prefix}typing.Union`\\[{{}}]]"
-            args = tuple(x for x in args if x is not type(None))  # noqa: E721
+        else:
+            simplify_optional_unions: bool = getattr(config, "simplify_optional_unions", True)
+            if not simplify_optional_unions:
+                full_name = "typing.Optional"
+                args_format = f"\\[:py:data:`{prefix}typing.Union`\\[{{}}]]"
+                args = tuple(x for x in args if x is not type(None))  # noqa: E721
     elif full_name == "typing.Callable" and args and args[0] is not ...:
-        fmt = ", ".join(format_annotation(arg, simplify_optional_unions=simplify_optional_unions) for arg in args[:-1])
-        formatted_args = f"\\[\\[{fmt}]"
-        formatted_args += f", {format_annotation(args[-1], simplify_optional_unions=simplify_optional_unions)}]"
+        fmt = [format_annotation(arg, config) for arg in args]
+        formatted_args = f"\\[\\[{', '.join(fmt[:-1])}], {fmt[-1]}]"
     elif full_name == "typing.Literal":
         formatted_args = f"\\[{', '.join(repr(arg) for arg in args)}]"
 
     if args and not formatted_args:
-        fmt = ", ".join(format_annotation(arg, fully_qualified, simplify_optional_unions) for arg in args)
-        formatted_args = args_format.format(fmt)
+        fmt = [format_annotation(arg, config) for arg in args]
+        formatted_args = args_format.format(", ".join(fmt))
 
     return f":py:{role}:`{prefix}{full_name}`{formatted_args}"
 
@@ -438,11 +446,6 @@ def process_docstring(
         signature = None
     type_hints = get_all_type_hints(obj, name)
 
-    formatter = partial(
-        format_annotation,
-        fully_qualified=app.config.typehints_fully_qualified,
-        simplify_optional_unions=app.config.simplify_optional_unions,
-    )
     for arg_name, annotation in type_hints.items():
         if arg_name == "return":
             continue  # this is handled separately later
@@ -453,7 +456,7 @@ def process_docstring(
         if arg_name.endswith("_"):
             arg_name = f"{arg_name[:-1]}\\_"
 
-        formatted_annotation = formatter(annotation)
+        formatted_annotation = format_annotation(annotation, app.config)
 
         search_for = {f":{field} {arg_name}:" for field in ("param", "parameter", "arg", "argument")}
         insert_index = None
@@ -480,7 +483,7 @@ def process_docstring(
     if "return" in type_hints and not inspect.isclass(original_obj):
         if what == "method" and name.endswith(".__init__"):  # avoid adding a return type for data class __init__
             return
-        formatted_annotation = formatter(type_hints["return"])
+        formatted_annotation = format_annotation(type_hints["return"], app.config)
         insert_index = len(lines)
         for at, line in enumerate(lines):
             if line.startswith(":rtype:"):
@@ -506,6 +509,10 @@ def validate_config(app: Sphinx, env: BuildEnvironment, docnames: list[str]) ->
     if app.config.typehints_defaults not in valid | {False}:
         raise ValueError(f"typehints_defaults needs to be one of {valid!r}, not {app.config.typehints_defaults!r}")
 
+    formatter = app.config.typehints_formatter
+    if formatter is not None and not callable(formatter):
+        raise ValueError(f"typehints_formatter needs to be callable or `None`, not {formatter}")
+
 
 def setup(app: Sphinx) -> dict[str, bool]:
     app.add_config_value("set_type_checking_flag", False, "html")
@@ -514,6 +521,7 @@ def setup(app: Sphinx) -> dict[str, bool]:
     app.add_config_value("typehints_document_rtype", True, "env")
     app.add_config_value("typehints_defaults", None, "env")
     app.add_config_value("simplify_optional_unions", True, "env")
+    app.add_config_value("typehints_formatter", None, "env")
     app.connect("builder-inited", builder_ready)
     app.connect("env-before-read-docs", validate_config)  # config may be changed after “config-inited” event
     app.connect("autodoc-process-signature", process_signature)

tests/test_sphinx_autodoc_typehints.py

@@ -203,7 +203,8 @@ def test_parse_annotation(annotation: Any, module: str, class_name: str, args: t
     ],
 )
 def test_format_annotation(inv: Inventory, annotation: Any, expected_result: str) -> None:
-    result = format_annotation(annotation)
+    conf = create_autospec(Config)
+    result = format_annotation(annotation, conf)
     assert result == expected_result
 
     # Test with the "simplify_optional_unions" flag turned off:
@@ -214,21 +215,21 @@ def test_format_annotation(inv: Inventory, annotation: Any, expected_result: str
         # encapsulate Union in typing.Optional
         expected_result_not_simplified = ":py:data:`~typing.Optional`\\[" + expected_result_not_simplified
         expected_result_not_simplified += "]"
-        assert format_annotation(annotation, simplify_optional_unions=False) == expected_result_not_simplified
+        conf = create_autospec(Config, simplify_optional_unions=False)
+        assert format_annotation(annotation, conf) == expected_result_not_simplified
 
         # Test with the "fully_qualified" flag turned on
         if "typing" in expected_result_not_simplified:
             expected_result_not_simplified = expected_result_not_simplified.replace("~typing", "typing")
-            assert (
-                format_annotation(annotation, fully_qualified=True, simplify_optional_unions=False)
-                == expected_result_not_simplified
-            )
+            conf = create_autospec(Config, fully_qualified=True, simplify_optional_unions=False)
+            assert format_annotation(annotation, conf) == expected_result_not_simplified
 
     # Test with the "fully_qualified" flag turned on
     if "typing" in expected_result or __name__ in expected_result:
         expected_result = expected_result.replace("~typing", "typing")
         expected_result = expected_result.replace("~" + __name__, __name__)
-        assert format_annotation(annotation, fully_qualified=True) == expected_result
+        conf = create_autospec(Config, fully_qualified=True)
+        assert format_annotation(annotation, conf) == expected_result
 
     # Test for the correct role (class vs data) using the official Sphinx inventory
     if "typing" in expected_result:
@@ -262,13 +263,15 @@ def test_format_annotation_both_libs(library: ModuleType, annotation: str, param
         return  # pragma: no cover
 
     ann = annotation_cls if params is None else annotation_cls[params]
-    result = format_annotation(ann)
+    result = format_annotation(ann, create_autospec(Config))
     assert result == expected_result
 
 
 def test_process_docstring_slot_wrapper() -> None:
     lines: list[str] = []
-    config = create_autospec(Config, typehints_fully_qualified=False, simplify_optional_unions=False)
+    config = create_autospec(
+        Config, typehints_fully_qualified=False, simplify_optional_unions=False, typehints_formatter=None
+    )
     app: Sphinx = create_autospec(Sphinx, config=config)
     process_docstring(app, "class", "SlotWrapper", Slotted, None, lines)
     assert not lines
@@ -679,6 +682,54 @@ def test_sphinx_output_defaults(
     assert text_contents == dedent(expected_contents)
 
 
+@pytest.mark.parametrize(
+    ("formatter_config_val", "expected"),
+    [
+        (None, ['("bool") -- foo', '("int") -- bar', '"str"']),
+        (lambda ann, conf: "Test", ["(*Test*) -- foo", "(*Test*) -- bar", "Test"]),
+        ("some string", Exception("needs to be callable or `None`")),
+    ],
+)
+@pytest.mark.sphinx("text", testroot="dummy")
+@patch("sphinx.writers.text.MAXWIDTH", 2000)
+def test_sphinx_output_formatter(
+    app: SphinxTestApp, status: StringIO, formatter_config_val: str, expected: tuple[str, ...] | Exception
+) -> None:
+    set_python_path()
+
+    app.config.master_doc = "simple"  # type: ignore # create flag
+    app.config.typehints_formatter = formatter_config_val  # type: ignore # create flag
+    try:
+        app.build()
+    except Exception as e:
+        if not isinstance(expected, Exception):
+            raise
+        assert str(expected) in str(e)
+        return
+    assert not isinstance(expected, Exception), "Expected app.build() to raise exception, but it didn’t"
+    assert "build succeeded" in status.getvalue()
+
+    text_path = pathlib.Path(app.srcdir) / "_build" / "text" / "simple.txt"
+    text_contents = text_path.read_text().replace("–", "--")
+    expected_contents = f"""\
+    Simple Module
+    *************
+
+    dummy_module_simple.function(x, y=1)
+
+       Function docstring.
+
+       Parameters:
+          * **x** {expected[0]}
+
+          * **y** {expected[1]}
+
+       Return type:
+          {expected[2]}
+    """
+    assert text_contents == dedent(expected_contents)
+
+
 def test_normalize_source_lines_async_def() -> None:
     source = """
     async def async_function():
@@ -733,7 +784,9 @@ def __init__(bound_args):  # noqa: N805
 
 @pytest.mark.parametrize("obj", [cmp_to_key, 1])
 def test_default_no_signature(obj: Any) -> None:
-    config = create_autospec(Config, typehints_fully_qualified=False, simplify_optional_unions=False)
+    config = create_autospec(
+        Config, typehints_fully_qualified=False, simplify_optional_unions=False, typehints_formatter=None
+    )
     app: Sphinx = create_autospec(Sphinx, config=config)
     lines: list[str] = []
     process_docstring(app, "what", "name", obj, None, lines)
@@ -749,6 +802,7 @@ def test_bound_class_method(method: FunctionType) -> None:
         typehints_document_rtype=False,
         always_document_param_types=True,
         typehints_defaults=True,
+        typehints_formatter=None,
     )
     app: Sphinx = create_autospec(Sphinx, config=config)
     process_docstring(app, "class", method.__qualname__, method, None, [])

whitelist.txt

@@ -2,6 +2,7 @@ ast3
 autodoc
 autouse
 backfill
+conf
 contravariant
 cpython
 dedent