From bfaf30dbcc4a3eb49f14c93f9be9d8a42f069ba6 Mon Sep 17 00:00:00 2001
From: tm91236 <205063739+tm91236@users.noreply.github.com>
Date: Mon, 31 Mar 2025 12:11:26 +0000
Subject: [PATCH] fix(docs): remove `typehints_formatter`.

By removing the `typehints_formatter` we are able to resolve
issues #570 and #734. The former issue is resolved by removing
the non-picklable function, which enables build caching, and the
second is solved by removing the `TypeVar` bound replacement.

While the `TypeVar` replacement does provide more concise
documentation, it removes a key piece of information, which the
reader must be aware of. For example, in:
```python
T = TypeVar("T", bound=Data)

def typevar(a: T, b: T) -> None: ...
def bound(a: Data, b: Data): ...

def valid(a: SupervisedData, b: SupervisedData): ...
def invalid(a: Data, b: SupervisedData): ...
```
the `typevar` function is the actual specification, `bound` is the
documented specification, `valid` is a valid specification, and
`invalid` is an invalid specification which is presented as valid
in the documentation (if we replace the `TypeVar` with the bound).

As an additional comment, we do now lose the functionality of the
`autodoc_custom_types` -- it should be noted that the `ArrayLike`
custom type (the only listed custom type) doesn't appear to be
used anywhere in the codebase.

Refs: #570, #734
---
 documentation/source/conf.py | 47 +-----------------------------------
 1 file changed, 1 insertion(+), 46 deletions(-)

diff --git a/documentation/source/conf.py b/documentation/source/conf.py
index e41c8503..72a9836a 100644
--- a/documentation/source/conf.py
+++ b/documentation/source/conf.py
@@ -23,14 +23,11 @@
 import sys
 from pathlib import Path
 from types import ModuleType
-from typing import Any, Optional, TypeVar, Union
+from typing import Optional
 from unittest import mock
 
-import sphinx.config
 import sphobjinv
 import tqdm
-from jax.typing import ArrayLike
-from sphinx_autodoc_typehints import format_annotation as default_format_annotation
 
 # https://docs.github.com/en/actions/learn-github-actions/variables,
 # see the "Default environment variables" section
@@ -78,11 +75,6 @@
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration
 
-# pylint: disable=invalid-name
-# Fixes ISSUE #561 pending a fix for ISSUE #570
-show_warning_types = True
-suppress_warnings = ["config.cache"]
-# pylint: enable=invalid-name
 
 # sphinx extensions
 extensions = [
@@ -216,10 +208,6 @@
     ("py:class", r"jaxtyping\.Shaped\[Array, '.*']"),
 ]
 
-autodoc_custom_types: dict[Any, str] = {  # Specify custom types for autodoc_type_hints
-    ArrayLike: ":data:`~jax.typing.ArrayLike`",
-}
-
 # custom references for tqdm, which does not support intersphinx
 tqdm_refs: dict[str, dict[str, str]] = {
     "py:class": {
@@ -227,39 +215,6 @@
     }
 }
 
-
-def typehints_formatter(
-    annotation: Any, config: sphinx.config.Config
-) -> Union[str, None]:
-    """
-    Properly replace custom type aliases.
-
-    :param annotation: The type annotation to be processed.
-    :param config: The current configuration being used.
-    :returns: A string of reStructured text (e.g. :py:class:`something`) or None to fall
-        back to the default.
-
-    This function is called on each type annotation that Sphinx processes.
-    The following steps occur:
-
-    1. Check if the annotation is a TypeVar. If so, replace it with its "bound" type
-        for clarity in the docs. If not, then replace it with typing.Any.
-    2. Check whether a specific Sphinx string has been defined in autodoc_custom_types.
-        If so, return that.
-    3. If not, then return None, which uses thee default formatter.
-
-    See https://github.com/tox-dev/sphinx-autodoc-typehints?tab=readme-ov-file#options
-    for specification.
-    """
-    if isinstance(annotation, TypeVar):
-        if annotation.__bound__ is None:  # when a generic TypeVar has been used.
-            return default_format_annotation(Any, config)
-        return default_format_annotation(
-            annotation.__bound__, config
-        )  # get the annotation for the bound type
-    return autodoc_custom_types.get(annotation)
-
-
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output