diff --git a/docs/source/_static/overload_dropdown.css b/docs/source/_static/overload_dropdown.css
new file mode 100644
index 00000000..b62ca980
--- /dev/null
+++ b/docs/source/_static/overload_dropdown.css
@@ -0,0 +1,53 @@
+/* ── Toggle button injected into the function-signature
────────────── */
+
+.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 , 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);
+}
diff --git a/docs/source/_static/overload_dropdown.js b/docs/source/_static/overload_dropdown.js
new file mode 100644
index 00000000..296b9f57
--- /dev/null
+++ b/docs/source/_static/overload_dropdown.js
@@ -0,0 +1,61 @@
+/**
+ * For each injected `.overloads-block` panel, find the function-signature
+ * `` 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 `` 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 , which is a child of .
+ var dl = block.closest("dl.py");
+ if (!dl) return;
+
+ // Use :scope to get a *direct* child - (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
- 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);
+ });
+ });
+});
diff --git a/docs/source/conf.py b/docs/source/conf.py
index ecad14af..9b004d75 100644
--- a/docs/source/conf.py
+++ b/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
diff --git a/docs/source/overload_dropdown.py b/docs/source/overload_dropdown.py
new file mode 100644
index 00000000..dd2a36a2
--- /dev/null
+++ b/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
- (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}