script_pages.py

"""Sphinx extension: generate per-script RST pages, copy .py into built _sources, and rewrite source links."""
import os
import glob
import io
from sphinx.util import logging

logger = logging.getLogger(__name__)

MARKER_TEMPLATE = '.. AUTO-GENERATED from {}'


def _title_from_script(path):
    """Extract a title from the first commented heading line; fallback to filename."""
    title = os.path.splitext(os.path.basename(path))[0]
    try:
        with io.open(path, 'r', encoding='utf-8') as fh:
            first_line = fh.readline()
            # Skip shebang (#!) if present
            if first_line and first_line.lstrip().startswith('#!'):
                first_line = fh.readline()
            # Look for first comment line
            title_comment = None
            if first_line and first_line.lstrip().startswith('#'):
                title_comment = first_line.strip().lstrip('#').strip()
            if title_comment:
                title = title_comment
    except Exception:
        pass
    return title


def _generate_rst_for_scripts(app):
    """Create a <name>.rst next to each <name>.py under docs/usage/examples/.

    - If `<name>.rst` already exists and does NOT start with our auto-generated marker,
      we leave it alone (so user-written RST isn't overwritten).
    - Otherwise we create or overwrite the RST to include the script via literalinclude.
    """
    logger = logging.getLogger(__name__)
    srcdir = app.srcdir

    patterns = [
        os.path.join(srcdir, 'usage', 'examples', '*.py'),
    ]
    py_files = []
    for pat in patterns:
        py_files.extend(glob.glob(pat))
    py_files = sorted(set(py_files))

    generated = 0

    for script_fpath in py_files:
        script_basename = os.path.basename(script_fpath)
        # Skip conf.py and other build files
        if script_basename in ('conf.py',):
            continue

        rst_path = os.path.splitext(script_fpath)[0] + '.rst'
        marker = MARKER_TEMPLATE.format(script_basename)

        # If rst exists and not generated by us, skip it
        if os.path.exists(rst_path):
            try:
                with io.open(rst_path, 'r', encoding='utf-8') as rh:
                    first = rh.readline().strip()
                if first != marker:
                    logger.info(f'Skipping existing RST: {rst_path}')
                    continue
            except Exception:
                # If we can't read, skip to be safe
                logger.warning(f'Could not read existing RST {rst_path}; skipping')
                continue

        # Build content
        title = _title_from_script(script_fpath)
        underline = '=' * len(title)
        # Make the literalinclude path relative to the generated RST location
        rst_dir = os.path.dirname(rst_path)
        rel = os.path.relpath(script_fpath, rst_dir)
        lines = [
            marker,
            '',
            title,
            underline,
            '',
            f'.. literalinclude:: {rel}',
            '   :language: python',
            '   :linenos:',
            '',
        ]

        try:
            with io.open(rst_path, 'w', encoding='utf-8') as oh:
                oh.write('\n'.join(lines))
            generated += 1
            logger.info(f'Generated {rst_path} from {script_basename}')
        except Exception as e:
            logger.warning(f'Failed to write {rst_path}: {e}')

    logger.info(f'Generated {generated} script RST files')


def _remove_generated_rst(app, exception):
    """Remove RST files that were auto-generated for scripts and copy .py into built _sources.

    Safety rules:
    - Only remove files whose first line starts with the auto-generated marker
      ('.. AUTO-GENERATED from ...'). This avoids deleting user-authored RST.
    - Only remove files after a successful build (exception is None).
    """
    logger = logging.getLogger(__name__)

    # If build failed, do not remove files so devs can inspect outputs
    if exception is not None:
        logger.info('Build failed; leaving auto-generated RST files in place')
        return

    srcdir = app.srcdir
    patterns = [
        os.path.join(srcdir, 'usage', 'examples', '*.rst'),
    ]

    rst_files = []
    for pat in patterns:
        rst_files.extend(glob.glob(pat))

    removed = 0
    for rst in sorted(set(rst_files)):
        try:
            with io.open(rst, 'r', encoding='utf-8') as fh:
                first = fh.readline().strip()
            if first.startswith(MARKER_TEMPLATE.format('')):
                # If building HTML, copy the original .py into the built _sources
                # directory so the "View source" link can serve the real Python file.
                try:
                    py_full = os.path.splitext(rst)[0] + '.py'
                    if os.path.exists(py_full) and getattr(app, 'builder', None) is not None:
                        try:
                            builder_name = getattr(app.builder, 'name', None)
                            outdir = getattr(app.builder, 'outdir', None)
                        except Exception:
                            builder_name = None
                            outdir = None

                        if builder_name == 'html' and outdir:
                            dest_dir = os.path.join(outdir, '_sources')
                            rel_py = os.path.relpath(py_full, app.srcdir)
                            dest_path = os.path.join(dest_dir, rel_py)
                            os.makedirs(os.path.dirname(dest_path), exist_ok=True)
                            # Copy file contents
                            try:
                                with io.open(py_full, 'r', encoding='utf-8') as rf, io.open(dest_path, 'w', encoding='utf-8') as wf:
                                    wf.write(rf.read())
                                logger.info(f'Copied Python source to built _sources: {dest_path}')
                            except Exception as e:
                                logger.warning(f'Failed to write built source {dest_path}: {e}')

                except Exception:
                    # non-fatal; continue to removal attempt
                    pass

                try:
                    os.remove(rst)
                    removed += 1
                    logger.info(f'Removed generated RST: {rst}')
                except Exception as e:
                    logger.warning(f'Failed to remove generated RST {rst}: {e}')
        except Exception as e:
            logger.warning(f'Could not read {rst}; skipping removal: {e}')

    logger.info(f'Removed {removed} auto-generated script RST files')


def _rewrite_sourcelink_to_py(app, pagename, templatename, context, doctree):
    """If the current page was generated from a script RST, point "View source" to the .py.

    This mirrors nbsphinx behaviour for notebooks where the "Show source" link
    points to the original notebook instead of the generated RST.
    """
    try:
        # doc2path with base=None returns a path relative to srcdir
        rst_rel = app.env.doc2path(pagename, base=None)
    except Exception:
        return

    rst_full = os.path.join(app.srcdir, rst_rel)
    if not os.path.exists(rst_full):
        return

    try:
        with io.open(rst_full, 'r', encoding='utf-8') as fh:
            first = fh.readline().strip()
    except Exception:
        return

    # Only rewrite for our auto-generated script RST files
    if not first.startswith(MARKER_TEMPLATE.format('')):
        return

    py_full = os.path.splitext(rst_full)[0] + '.py'
    if not os.path.exists(py_full):
        return

    # Make the path relative to the source directory (what Sphinx expects)
    rel_py = os.path.relpath(py_full, app.srcdir)

    # Replace the template context variable so the HTML theme will link to .py
    context['sourcename'] = rel_py


def setup(app):
    # Generate per-script RST files before reading sources
    app.connect('builder-inited', _generate_rst_for_scripts)
    # Remove the generated per-script RST files after a successful build
    app.connect('build-finished', _remove_generated_rst)
    # Make "View source" point to the original .py for auto-generated script pages
    app.connect('html-page-context', _rewrite_sourcelink_to_py)

    return {
        'version': '1.0',
        'parallel_read_safe': True,
        'parallel_write_safe': True,
    }