docs: Add collapsible overload dropdown to function signatures

docs/source/_static/overload_dropdown.css

@@ -0,0 +1,53 @@
+/* ── Toggle button injected into the function-signature <dt> ────────────── */
+
+.overloads-toggle {
+    float: right;
+    display: inline-flex;
+    align-items: center;
+    gap: 0.3em;
+    margin-left: 1em;
+    padding: 0.1em 0.5em;
+    font-size: 0.75em;
+    font-family: var(--font-stack);
+    color: var(--color-foreground-secondary);
+    background: transparent;
+    border: 1px solid var(--color-foreground-muted, currentColor);
+    border-radius: 0.25em;
+    cursor: pointer;
+    opacity: 0.7;
+    user-select: none;
+    line-height: 1.8;
+    vertical-align: middle;
+}
+
+.overloads-toggle:hover {
+    opacity: 1;
+    color: var(--color-foreground-primary);
+}
+
+/* Animated chevron symbol supplied via CSS so it can be rotated smoothly. */
+.overloads-toggle-chevron::before {
+    content: "▶";
+    display: inline-block;
+    font-size: 0.7em;
+    transition: transform 0.2s ease;
+}
+
+.overloads-toggle--open .overloads-toggle-chevron::before {
+    transform: rotate(90deg);
+}
+
+/* When the button precedes a [source] link in the same <dt>, push [source]
+   below the button by clearing its float.  This only activates for functions
+   that have overloads (where the button is present). */
+.overloads-toggle ~ a.reference .viewcode-link {
+    clear: right;
+}
+
+/* ── Overloads content panel ─────────────────────────────────────────────── */
+
+.overloads-block {
+    margin-bottom: 1em;
+    padding-left: 1em;
+    border-left: 2px solid var(--color-brand-primary);
+}

docs/source/_static/overload_dropdown.js

@@ -0,0 +1,61 @@
+/**
+ * For each injected `.overloads-block` panel, find the function-signature
+ * `<dt>` that owns it and insert an "Overloads ▶" toggle button into it.
+ *
+ * Layout strategy
+ * ───────────────
+ * Both the button and the existing "[source]" link use `float: right`.
+ * Inserting the button *before* "[source]" in the DOM, combined with
+ * `clear: right` on the "[source]" float (applied via CSS), causes them to
+ * stack vertically rather than sit side by side.
+ *
+ * `display: flow-root` is set on `<dt>` so it establishes a block formatting
+ * context and expands to contain both floated elements.  Functions that have
+ * no overloads are never touched.
+ */
+document.addEventListener("DOMContentLoaded", function () {
+  document.querySelectorAll(".overloads-block").forEach(function (block) {
+    // The block lives inside <dd>, which is a child of <dl class="py …">.
+    var dl = block.closest("dl.py");
+    if (!dl) return;
+
+    // Use :scope to get a *direct* child <dt> (avoids matching nested dls).
+    var dt = dl.querySelector(":scope > dt");
+    if (!dt) return;
+
+    // Start collapsed.
+    block.hidden = true;
+
+    // Build the button.
+    var btn = document.createElement("button");
+    btn.type = "button";
+    btn.className = "overloads-toggle";
+    btn.setAttribute("aria-expanded", "false");
+
+    var label = document.createElement("span");
+    label.textContent = "Overloads";
+
+    var chevron = document.createElement("span");
+    chevron.className = "overloads-toggle-chevron";
+    chevron.setAttribute("aria-hidden", "true");
+
+    btn.appendChild(label);
+    btn.appendChild(chevron);
+
+    // Insert immediately before "[source]" so CSS `clear: right` on the
+    // viewcode-link makes it wrap below the button.
+    // When there is no "[source]" link, append at the end instead.
+    var sourceLink = dt.querySelector(".viewcode-link")?.closest("a");
+    dt.insertBefore(btn, sourceLink ?? null);
+
+    // Make <dt> contain both floated elements (button + [source]).
+    dt.style.display = "flow-root";
+
+    btn.addEventListener("click", function () {
+      var opening = block.hidden;
+      block.hidden = !opening;
+      btn.setAttribute("aria-expanded", opening ? "true" : "false");
+      btn.classList.toggle("overloads-toggle--open", opening);
+    });
+  });
+});

docs/source/conf.py

@@ -16,6 +16,8 @@
 _PATH_ROOT = os.path.realpath(os.path.join(_PATH_HERE, "..", ".."))
 _PATH_PYPROJECT = os.path.join(_PATH_ROOT, "pyproject.toml")
 
+sys.path.insert(0, _PATH_HERE)
+
 # Read all metadata from pyproject.toml, so that we don't duplicate it.
 
 with open(_PATH_PYPROJECT, mode="rb") as fp:
@@ -40,6 +42,7 @@
     "sphinx.ext.intersphinx",
     "myst_parser",  # Enables markdown support
     "sphinx_design",  # Enables side to side cards
+    "overload_dropdown",  # Shows @overload stubs in a collapsible dropdown
 ]
 
 # -- Options for HTML output -------------------------------------------------
@@ -69,8 +72,10 @@
     "sidebar_hide_name": True,
 }
 
+html_css_files = ["overload_dropdown.css"]
 html_js_files = [
     ("https://stats.torchjd.org/js/script.js", {"data-domain": "torchjd.org", "defer": "defer"}),
+    "overload_dropdown.js",
 ]
 
 html_title = "TorchJD"
@@ -108,21 +113,22 @@ def _get_obj(_info: dict[str, str]) -> object:
     for part in full_name.split("."):
         obj = getattr(obj, part)
     # strip decorators, which would resolve to the source of the decorator
-    obj = inspect.unwrap(obj)
+    obj = inspect.unwrap(obj)  # type: ignore[arg-type]
     return obj
 
 
 def _get_file_name(obj: object) -> str | None:
     try:
-        file_name = inspect.getsourcefile(obj)
-        file_name = os.path.relpath(file_name, start=_PATH_ROOT)
+        file_name = inspect.getsourcefile(obj)  # type: ignore[arg-type]
+        if file_name is None:
+            return None
+        return os.path.relpath(file_name, start=_PATH_ROOT)
     except TypeError:  # This seems to happen when obj is a property
-        file_name = None
-    return file_name
+        return None
 
 
 def _get_line_str(obj: object) -> str:
-    source, start = inspect.getsourcelines(obj)
+    source, start = inspect.getsourcelines(obj)  # type: ignore[arg-type]
     end = start + len(source) - 1
     line_str = f"#L{start}-L{end}"
     return line_str

docs/source/overload_dropdown.py

@@ -0,0 +1,168 @@
+"""Sphinx extension that shows ``@overload`` stubs in a collapsible dropdown.
+
+For every function or method whose docstring is processed by autodoc, this
+extension calls :func:`typing.get_overloads` to retrieve the registered
+overload stubs and prepends a ``sphinx_design`` ``.. dropdown::`` block to
+the docstring so readers can inspect each overload signature without cluttering
+the main documentation page.
+
+Requires Python 3.11+ (for :func:`typing.get_overloads`) and the
+``sphinx-design`` Sphinx extension.
+"""
+
+import inspect
+import sys
+import types
+import typing
+from collections.abc import Sequence
+
+from sphinx.application import Sphinx
+
+
+def _format_annotation(ann: object) -> str:
+    """Format a type annotation using short (unqualified) class names."""
+    if ann is None or ann is type(None):
+        return "None"
+
+    # X | Y union types (Python 3.10+ syntax)
+    if isinstance(ann, types.UnionType):
+        return " | ".join(_format_annotation(a) for a in ann.__args__)
+
+    origin = getattr(ann, "__origin__", None)
+    if origin is not None:
+        args: Sequence[object] = getattr(ann, "__args__", None) or ()
+
+        if origin is typing.Union:
+            return " | ".join(_format_annotation(a) for a in args)
+
+        origin_name: str = (
+            getattr(origin, "__name__", None) or getattr(origin, "_name", None) or repr(origin)
+        )
+        if args:
+            args_str = ", ".join(_format_annotation(a) for a in args)
+            return f"{origin_name}[{args_str}]"
+        return origin_name
+
+    name: str | None = getattr(ann, "__name__", None)
+    if name:
+        return name
+
+    return str(ann)
+
+
+def _format_param(param: inspect.Parameter) -> str:
+    """Format a single parameter without separators (``/`` or ``*``)."""
+    if param.kind == inspect.Parameter.VAR_POSITIONAL:
+        name_part = f"*{param.name}"
+    elif param.kind == inspect.Parameter.VAR_KEYWORD:
+        name_part = f"**{param.name}"
+    else:
+        name_part = param.name
+
+    if param.annotation is not inspect.Parameter.empty:
+        name_part = f"{name_part}: {_format_annotation(param.annotation)}"
+
+    if param.default is not inspect.Parameter.empty:
+        name_part = f"{name_part} = {param.default!r}"
+
+    return name_part
+
+
+def _build_param_strs(sig: inspect.Signature) -> tuple[list[str], str]:
+    """Return ``(param_strs, return_str)`` for *sig* with short type names.
+
+    *param_strs* includes the bare ``/`` and ``*`` separators as individual
+    entries so that callers can join them however they like.
+    """
+    items = [(p.kind, _format_param(p)) for p in sig.parameters.values()]
+
+    param_strs: list[str] = []
+    for i, (kind, s) in enumerate(items):
+        prev_kind = items[i - 1][0] if i > 0 else None
+
+        # Insert '/' after the last positional-only parameter.
+        if (
+            prev_kind == inspect.Parameter.POSITIONAL_ONLY
+            and kind != inspect.Parameter.POSITIONAL_ONLY
+        ):
+            param_strs.append("/")
+
+        # Insert bare '*' before the first keyword-only parameter when there
+        # is no preceding *args parameter.
+        if kind == inspect.Parameter.KEYWORD_ONLY and prev_kind not in (
+            inspect.Parameter.KEYWORD_ONLY,
+            inspect.Parameter.VAR_POSITIONAL,
+        ):
+            param_strs.append("*")
+
+        param_strs.append(s)
+
+    # Trailing '/' when every parameter is positional-only.
+    if items and items[-1][0] == inspect.Parameter.POSITIONAL_ONLY:
+        param_strs.append("/")
+
+    ret = sig.return_annotation
+    return_str = f" -> {_format_annotation(ret)}" if ret is not inspect.Signature.empty else ""
+
+    return param_strs, return_str
+
+
+def _overload_code_lines(func_name: str, overload_func: object, indent: str) -> list[str]:
+    """Return source-style lines for one overload, indented by *indent*."""
+    try:
+        sig = inspect.signature(overload_func)  # type: ignore[arg-type]
+    except (ValueError, TypeError):
+        return [f"{indent}@overload", f"{indent}def {func_name}(...): ...", ""]
+
+    param_strs, return_str = _build_param_strs(sig)
+
+    lines = [
+        f"{indent}@overload",
+        f"{indent}def {func_name}(",
+        *[f"{indent}    {p}," for p in param_strs],
+        f"{indent}){return_str}: ...",
+    ]
+    lines.append("")
+    return lines
+
+
+def _process_docstring(
+    _app: Sphinx,
+    what: str,
+    name: str,
+    obj: object,
+    _options: object,
+    lines: list[str],
+) -> None:
+    """Prepend an *Overloads* dropdown to docstrings of overloaded functions."""
+    if what not in ("function", "method"):
+        return
+    if sys.version_info < (3, 11):
+        return
+
+    try:
+        overloads = typing.get_overloads(obj)
+    except Exception:
+        return
+
+    if not overloads:
+        return
+
+    func_name = name.split(".")[-1]
+
+    # A plain container div is used instead of a sphinx_design dropdown so
+    # that the toggle lives in the function-signature <dt> (added by JS) rather
+    # than as a separate collapsible card below it.
+    # Indentation: container content → 3 spaces; code-block content → 6 spaces.
+    dropdown: list[str] = [".. container:: overloads-block", "", "   .. code-block:: python", ""]
+    for overload_func in overloads:
+        dropdown.extend(_overload_code_lines(func_name, overload_func, indent="      "))
+    dropdown.append("")
+
+    for i, line in enumerate(dropdown):
+        lines.insert(i, line)
+
+
+def setup(app: Sphinx) -> dict[str, object]:
+    app.connect("autodoc-process-docstring", _process_docstring)
+    return {"version": "0.1", "parallel_read_safe": True}