docstring_renderer.py

"""Render braindecode rST docstrings to HTML for the model-explorer Space.

Braindecode docstrings use NumpyDoc + Sphinx extensions:
  .. figure::          (architecture image)
  :bdg-danger:         (sphinx-design colored badges)
  .. versionadded::    (Sphinx admonition)
  .. important::       (admonition)
  .. code-block::      (highlighted code)
  [ref]_               (NumpyDoc citation reference)

Pure docutils does not know about the Sphinx directives. Rather than
spinning up a full Sphinx build inside the Space, we:

1. Pre-process the docstring with regex substitutions that map
   Sphinx-only directives to plain rST equivalents docutils can parse.
2. Hand the result to docutils.publish_parts for rST -> HTML.
3. Wrap with a small CSS style block that matches braindecode.org colors.

This keeps the Space dependency-light (no sphinx at runtime) while still
rendering figures, headings, parameter tables, and references.
"""

from __future__ import annotations

import inspect
import re
from textwrap import dedent

from docutils.core import publish_parts
from docutils.utils import SystemMessage

_BADGE_COLORS = {
    "bdg-danger": "#d9534f",
    "bdg-success": "#5cb85c",
    "bdg-primary": "#0072B2",
    "bdg-info": "#56B4E9",
    "bdg-warning": "#E69F00",
    "bdg-secondary": "#6c757d",
    "bdg-light": "#f0f0f0",
    "bdg-dark": "#343a40",
}

_BADGE_RE = re.compile(r":(bdg-[a-z]+):`([^`]+)`")
_VERSIONADDED_RE = re.compile(r"^\.\. versionadded::\s*(.+)$", re.MULTILINE)
_VERSIONCHANGED_RE = re.compile(r"^\.\. versionchanged::\s*(.+)$", re.MULTILINE)
_CODE_BLOCK_RE = re.compile(r"^(\s*)\.\. code-block::\s*(\w+)?\s*$", re.MULTILINE)


def _replace_badges(text: str) -> str:
    """Convert :bdg-danger:`Foundation Model` to inline raw HTML."""

    def repl(match: re.Match) -> str:
        cls, label = match.group(1), match.group(2)
        color = _BADGE_COLORS.get(cls, "#888")
        # Use rST raw HTML inline pass-through.
        return (
            f"\n\n.. raw:: html\n\n"
            f"   <span style=\"display:inline-block;padding:2px 8px;"
            f"border-radius:4px;background:{color};color:white;"
            f"font-size:11px;font-weight:600;margin-right:4px;\">{label}</span>\n\n"
        )

    return _BADGE_RE.sub(repl, text)


def _replace_versionadded(text: str) -> str:
    """Convert Sphinx versionadded/versionchanged to plain admonitions."""
    text = _VERSIONADDED_RE.sub(
        r".. note::\n\n   *New in version \1.*", text
    )
    text = _VERSIONCHANGED_RE.sub(
        r".. note::\n\n   *Changed in version \1.*", text
    )
    return text


def _normalize_code_block(text: str) -> str:
    """Convert `.. code-block:: python` to vanilla `.. code::` (docutils-OK)."""

    def repl(match: re.Match) -> str:
        indent, lang = match.group(1), match.group(2) or ""
        return f"{indent}.. code:: {lang}".rstrip()

    return _CODE_BLOCK_RE.sub(repl, text)


def _strip_unsupported_directives(text: str) -> str:
    """Drop directives that docutils cannot parse and we do not want rendered.

    Currently: rubric (treated as a small heading) and bibliography-only items.
    """
    text = re.sub(r"^\.\. rubric::\s*(.+)$", r"\n**\1**\n", text, flags=re.MULTILINE)
    return text


def preprocess_docstring(doc: str) -> str:
    """Apply all rST → docutils-friendly transformations."""
    if not doc:
        return ""
    doc = dedent(doc)
    doc = _replace_badges(doc)
    doc = _replace_versionadded(doc)
    doc = _normalize_code_block(doc)
    doc = _strip_unsupported_directives(doc)
    return doc


_STYLE = """
<style>
  .bd-doc { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto,
            sans-serif; line-height: 1.55; color: #1f2937; }
  .bd-doc h1, .bd-doc h2, .bd-doc h3 { color: #0072B2; margin-top: 1.2em; }
  .bd-doc h1 { font-size: 1.5em; border-bottom: 2px solid #0072B2; padding-bottom: 4px; }
  .bd-doc h2 { font-size: 1.2em; }
  .bd-doc h3 { font-size: 1.05em; }
  .bd-doc pre { background: #f6f8fa; padding: 10px 12px; border-radius: 6px;
                font-size: 0.9em; overflow-x: auto; }
  .bd-doc code { background: #f0f0f5; padding: 1px 5px; border-radius: 3px;
                 font-size: 0.9em; }
  .bd-doc pre code { background: transparent; padding: 0; }
  .bd-doc img { max-width: 480px; display: block; margin: 12px auto;
                border-radius: 6px; }
  .bd-doc table { border-collapse: collapse; margin: 8px 0; }
  .bd-doc th, .bd-doc td { border: 1px solid #d0d7de; padding: 4px 10px;
                           text-align: left; }
  .bd-doc th { background: #f6f8fa; }
  .bd-doc .admonition { border-left: 4px solid #0072B2; background: #f0f7ff;
                        padding: 8px 14px; margin: 12px 0; border-radius: 4px; }
  .bd-doc .admonition.important { border-color: #D55E00; background: #fdf6ec; }
  .bd-doc .admonition.note { border-color: #009E73; background: #effaf3; }
  .bd-doc .admonition-title { font-weight: 600; margin-bottom: 4px; }
  .bd-doc dl.field-list { display: grid; grid-template-columns: max-content auto;
                          gap: 4px 12px; }
  .bd-doc dl.field-list dt { font-weight: 600; color: #475569; }
</style>
"""


def render_docstring_html(doc: str | None) -> str:
    """Render an rST docstring to a self-contained HTML fragment.

    Returns a string with embedded <style> + <div class="bd-doc">…</div>.
    Failures fall back to a <pre> dump so the Space never blanks out.
    """
    if not doc:
        return _STYLE + "<div class='bd-doc'><em>No docstring available.</em></div>"

    processed = preprocess_docstring(doc)
    try:
        parts = publish_parts(
            source=processed,
            writer_name="html5",
            settings_overrides={
                "report_level": 5,  # suppress all docutils warnings
                "halt_level": 5,
                "embed_stylesheet": False,
                "input_encoding": "unicode",
                "output_encoding": "unicode",
                "doctitle_xform": False,
                "initial_header_level": 2,
            },
        )
        body = parts["html_body"]
    except SystemMessage as exc:  # pragma: no cover — defensive
        body = f"<pre>{processed}</pre><p><em>(rST parse error: {exc})</em></p>"

    return _STYLE + f"<div class='bd-doc'>{body}</div>"


def get_signature_str(cls: type) -> str:
    """Return the formatted __init__ signature for display."""
    try:
        sig = inspect.signature(cls.__init__)
        return f"{cls.__name__}{sig}"
    except (ValueError, TypeError):
        return f"{cls.__name__}(...)"


def get_source_link(cls: type, branch: str = "master") -> str | None:
    """Return a github link to the class definition."""
    try:
        module = inspect.getmodule(cls)
        if module is None or not hasattr(module, "__file__"):
            return None
        rel_path = module.__file__.split("braindecode/", 1)[-1]
        rel_path = "braindecode/" + rel_path
        try:
            _, lineno = inspect.getsourcelines(cls)
        except OSError:
            lineno = 1
        return (
            f"https://github.com/braindecode/braindecode/blob/{branch}/"
            f"{rel_path}#L{lineno}"
        )
    except Exception:
        return None