Repository: tox-dev/sphinx-autodoc-typehints (python)

Issue:
Allow customizing formatter
Before making a PR, I’d like to check if you’d be open to the idea.

Basically, instead of `format_annotation(annotation, fully_qualified=False)`, we’d have

```py
class AnnotationFormatter:
    def __init__(self, *, fully_qualified=False):
        self.fully_qualified = fully_qualified
    def format(self, annotation) -> str:
        ...
```

Also a new config option `typehints_format_cls` where people would have the ability to override formatting:

```py
class CustomFormatter(sat.AnnotationFormatter):
    def format(self, annotation) -> str:
        if isinstance(annotation, MyClass):
            return ...
        return super().format(annotation)

typehints_format_cls = CustomFormatter
```

For backwards compatibility (if desirable), we could define

```py
def format_annotation(annotation, fully_qualified=False):
    return AnnotationFormatter(fully_qualified=fully_qualified).format(annotation)
```

Pull request:
Configurable formatter
Fixes #112

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index dc0e09b..af41039 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,7 +2,9 @@
 
 ## dev
 
+- 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
 
diff --git a/README.md b/README.md
index d894bb4..d03c3e9 100644
--- a/README.md
+++ b/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
 
diff --git a/setup.cfg b/setup.cfg
index a63428e..7d4e670 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -71,7 +71,7 @@ other =
     *\sphinx-autodoc-typehints
 
 [coverage:report]
-fail_under = 55
+fail_under = 78
 
 [coverage:html]
 show_contexts = true
diff --git a/src/sphinx_autodoc_typehints/__init__.py b/src/sphinx_autodoc_typehints/__init__.py
index ca5926d..e0f22e2 100644
--- a/src/sphinx_autodoc_typehints/__init__.py
+++ b/src/sphinx_autodoc_typehints/__init__.py
@@ -1,14 +1,15 @@
 from __future__ import annotations
 
 import inspect
+import re
 import sys
 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
@@ -89,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`"
@@ -115,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 = "\\[{}]"
@@ -130,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}"
 
@@ -252,43 +261,48 @@ def _future_annotations_imported(obj: Any) -> bool:
 
 
 def get_all_type_hints(obj: Any, name: str) -> dict[str, Any]:
-    rv = {}
-
+    result = _get_type_hint(name, obj)
+    if result:
+        return result
+    result = backfill_type_hints(obj, name)
     try:
-        rv = get_type_hints(obj)
-    except (AttributeError, TypeError, RecursionError) as exc:
-        # Introspecting a slot wrapper will raise TypeError, and and some recursive type
-        # definitions will cause a RecursionError (https://github.com/python/typing/issues/574).
+        obj.__annotations__ = result
+    except (AttributeError, TypeError):
+        return result
+    return _get_type_hint(name, obj)
 
-        # If one is using PEP563 annotations, Python will raise a (e.g.,)
-        # TypeError("TypeError: unsupported operand type(s) for |: 'type' and 'NoneType'")
-        # on 'str | None', therefore we accept TypeErrors with that error message
-        # if 'annotations' is imported from '__future__'.
-        if isinstance(exc, TypeError) and _future_annotations_imported(obj) and "unsupported operand type" in str(exc):
-            rv = obj.__annotations__
-    except NameError as exc:
-        _LOGGER.warning('Cannot resolve forward reference in type annotations of "%s": %s', name, exc)
-        rv = obj.__annotations__
 
-    if rv:
-        return rv
+_TYPE_GUARD_IMPORT_RE = re.compile(r"if (typing.)?TYPE_CHECKING:([\s\S]*?)(?=\n\S)")
+_TYPE_GUARD_IMPORTS_RESOLVED = set()
 
-    rv = backfill_type_hints(obj, name)
 
-    try:
-        obj.__annotations__ = rv
-    except (AttributeError, TypeError):
-        return rv
+def _resolve_type_guarded_imports(obj: Any) -> None:
+    if hasattr(obj, "__module__") and obj.__module__ not in _TYPE_GUARD_IMPORTS_RESOLVED:
+        _TYPE_GUARD_IMPORTS_RESOLVED.add(obj.__module__)
+        if obj.__module__ not in sys.builtin_module_names:
+            module = inspect.getmodule(obj)
+            if module:
+                module_code = inspect.getsource(module)
+                for (_, part) in _TYPE_GUARD_IMPORT_RE.findall(module_code):
+                    module_code = textwrap.dedent(part)
+                    exec(module_code, obj.__globals__)
+
 
+def _get_type_hint(name: str, obj: Any) -> dict[str, Any]:
+    _resolve_type_guarded_imports(obj)
     try:
-        rv = get_type_hints(obj)
-    except (AttributeError, TypeError):
-        pass
+        result = get_type_hints(obj)
+    except (AttributeError, TypeError, RecursionError) as exc:
+        # TypeError - slot wrapper, PEP-563 when part of new syntax not supported
+        # RecursionError - some recursive type definitions https://github.com/python/typing/issues/574
+        if isinstance(exc, TypeError) and _future_annotations_imported(obj) and "unsupported operand type" in str(exc):
+            result = obj.__annotations__
+        else:
+            result = {}
     except NameError as exc:
         _LOGGER.warning('Cannot resolve forward reference in type annotations of "%s": %s', name, exc)
-        rv = obj.__annotations__
-
-    return rv
+        result = obj.__annotations__
+    return result
 
 
 def backfill_type_hints(obj: Any, name: str) -> dict[str, Any]:
@@ -305,11 +319,9 @@ def backfill_type_hints(obj: Any, name: str) -> dict[str, Any]:
 
     def _one_child(module: Module) -> stmt | None:
         children = module.body  # use the body to ignore type comments
-
         if len(children) != 1:
             _LOGGER.warning('Did not get exactly one node from AST for "%s", got %s', name, len(children))
             return None
-
         return children[0]
 
     try:
@@ -434,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
@@ -449,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
@@ -476,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:"):
@@ -502,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")
@@ -510,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)
@@ -526,4 +538,5 @@ __all__ = [
     "normalize_source_lines",
     "process_docstring",
     "process_signature",
+    "backfill_type_hints",
 ]
diff --git a/whitelist.txt b/whitelist.txt
index 840c30c..6329db9 100644
--- a/whitelist.txt
+++ b/whitelist.txt
@@ -2,6 +2,7 @@ ast3
 autodoc
 autouse
 backfill
+conf
 contravariant
 cpython
 dedent
@@ -16,6 +17,7 @@ formatter
 func
 getmodule
 getsource
+globals
 idx
 inited
 inv
@@ -29,7 +31,6 @@ param
 parametrized
 params
 pathlib
-pep563
 pos
 prepend
 py310


Test patch:
diff --git a/tests/roots/test-resolve-typing-guard/conf.py b/tests/roots/test-resolve-typing-guard/conf.py
new file mode 100644
index 0000000..bb939ea
--- /dev/null
+++ b/tests/roots/test-resolve-typing-guard/conf.py
@@ -0,0 +1,9 @@
+import pathlib
+import sys
+
+master_doc = "index"
+sys.path.insert(0, str(pathlib.Path(__file__).parent))
+extensions = [
+    "sphinx.ext.autodoc",
+    "sphinx_autodoc_typehints",
+]
diff --git a/tests/roots/test-resolve-typing-guard/demo_typing_guard.py b/tests/roots/test-resolve-typing-guard/demo_typing_guard.py
new file mode 100644
index 0000000..03ad07e
--- /dev/null
+++ b/tests/roots/test-resolve-typing-guard/demo_typing_guard.py
@@ -0,0 +1,32 @@
+"""Module demonstrating imports that are type guarded"""
+from __future__ import annotations
+
+import typing
+from builtins import ValueError  # handle does not have __module__
+from functools import cmp_to_key  # has __module__ but cannot get module as is builtin
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from decimal import Decimal
+    from typing import Sequence
+
+if typing.TYPE_CHECKING:
+    from typing import AnyStr
+
+
+def a(f: Decimal, s: AnyStr) -> Sequence[AnyStr | Decimal]:
+    """
+    Do.
+
+    :param f: first
+    :param s: second
+    :return: result
+    """
+    return [f, s]
+
+
+__all__ = [
+    "a",
+    "ValueError",
+    "cmp_to_key",
+]
diff --git a/tests/roots/test-resolve-typing-guard/index.rst b/tests/roots/test-resolve-typing-guard/index.rst
new file mode 100644
index 0000000..f92eae8
--- /dev/null
+++ b/tests/roots/test-resolve-typing-guard/index.rst
@@ -0,0 +1,2 @@
+.. automodule:: demo_typing_guard
+   :members:
diff --git a/tests/test_sphinx_autodoc_typehints.py b/tests/test_sphinx_autodoc_typehints.py
index 13deb9b..e5a9b4c 100644
--- a/tests/test_sphinx_autodoc_typehints.py
+++ b/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 test_normalize_source_lines_def_starting_decorator_parameter() -> None:
 
 @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, [])
@@ -761,3 +815,11 @@ def test_syntax_error_backfill() -> None:
         lambda x: x)
     # fmt: on
     backfill_type_hints(func, "func")
+
+
+@pytest.mark.sphinx("text", testroot="resolve-typing-guard")
+def test_resolve_typing_guard_imports(app: SphinxTestApp, status: StringIO, warning: StringIO) -> None:
+    set_python_path()
+    app.build()
+    assert "build succeeded" in status.getvalue()
+    assert not warning.getvalue()