|
| 1 | +"""Sphinx extension: generate per-script RST pages, copy .py into built _sources, and rewrite source links.""" |
| 2 | +import os |
| 3 | +import glob |
| 4 | +import io |
| 5 | +from sphinx.util import logging |
| 6 | + |
| 7 | +logger = logging.getLogger(__name__) |
| 8 | + |
| 9 | +MARKER_TEMPLATE = '.. AUTO-GENERATED from {}' |
| 10 | + |
| 11 | + |
| 12 | +def _title_from_script(path): |
| 13 | + """Extract a title from the first commented heading line; fallback to filename.""" |
| 14 | + title = os.path.splitext(os.path.basename(path))[0] |
| 15 | + try: |
| 16 | + with io.open(path, 'r', encoding='utf-8') as fh: |
| 17 | + first_line = fh.readline() |
| 18 | + # Skip shebang (#!) if present |
| 19 | + if first_line and first_line.lstrip().startswith('#!'): |
| 20 | + first_line = fh.readline() |
| 21 | + # Look for first comment line |
| 22 | + title_comment = None |
| 23 | + if first_line and first_line.lstrip().startswith('#'): |
| 24 | + title_comment = first_line.strip().lstrip('#').strip() |
| 25 | + if title_comment: |
| 26 | + title = title_comment |
| 27 | + except Exception: |
| 28 | + pass |
| 29 | + return title |
| 30 | + |
| 31 | + |
| 32 | +def _generate_rst_for_scripts(app): |
| 33 | + """Create a <name>.rst next to each <name>.py under docs/usage/examples/. |
| 34 | +
|
| 35 | + - If `<name>.rst` already exists and does NOT start with our auto-generated marker, |
| 36 | + we leave it alone (so user-written RST isn't overwritten). |
| 37 | + - Otherwise we create or overwrite the RST to include the script via literalinclude. |
| 38 | + """ |
| 39 | + logger = logging.getLogger(__name__) |
| 40 | + srcdir = app.srcdir |
| 41 | + |
| 42 | + patterns = [ |
| 43 | + os.path.join(srcdir, 'usage', 'examples', '*.py'), |
| 44 | + ] |
| 45 | + py_files = [] |
| 46 | + for pat in patterns: |
| 47 | + py_files.extend(glob.glob(pat)) |
| 48 | + py_files = sorted(set(py_files)) |
| 49 | + |
| 50 | + generated = 0 |
| 51 | + |
| 52 | + for script_fpath in py_files: |
| 53 | + script_basename = os.path.basename(script_fpath) |
| 54 | + # Skip conf.py and other build files |
| 55 | + if script_basename in ('conf.py',): |
| 56 | + continue |
| 57 | + |
| 58 | + rst_path = os.path.splitext(script_fpath)[0] + '.rst' |
| 59 | + marker = MARKER_TEMPLATE.format(script_basename) |
| 60 | + |
| 61 | + # If rst exists and not generated by us, skip it |
| 62 | + if os.path.exists(rst_path): |
| 63 | + try: |
| 64 | + with io.open(rst_path, 'r', encoding='utf-8') as rh: |
| 65 | + first = rh.readline().strip() |
| 66 | + if first != marker: |
| 67 | + logger.info(f'Skipping existing RST: {rst_path}') |
| 68 | + continue |
| 69 | + except Exception: |
| 70 | + # If we can't read, skip to be safe |
| 71 | + logger.warning(f'Could not read existing RST {rst_path}; skipping') |
| 72 | + continue |
| 73 | + |
| 74 | + # Build content |
| 75 | + title = _title_from_script(script_fpath) |
| 76 | + underline = '=' * len(title) |
| 77 | + # Make the literalinclude path relative to the generated RST location |
| 78 | + rst_dir = os.path.dirname(rst_path) |
| 79 | + rel = os.path.relpath(script_fpath, rst_dir) |
| 80 | + lines = [ |
| 81 | + marker, |
| 82 | + '', |
| 83 | + title, |
| 84 | + underline, |
| 85 | + '', |
| 86 | + f'.. literalinclude:: {rel}', |
| 87 | + ' :language: python', |
| 88 | + ' :linenos:', |
| 89 | + '', |
| 90 | + ] |
| 91 | + |
| 92 | + try: |
| 93 | + with io.open(rst_path, 'w', encoding='utf-8') as oh: |
| 94 | + oh.write('\n'.join(lines)) |
| 95 | + generated += 1 |
| 96 | + logger.info(f'Generated {rst_path} from {script_basename}') |
| 97 | + except Exception as e: |
| 98 | + logger.warning(f'Failed to write {rst_path}: {e}') |
| 99 | + |
| 100 | + logger.info(f'Generated {generated} script RST files') |
| 101 | + |
| 102 | + |
| 103 | +def _remove_generated_rst(app, exception): |
| 104 | + """Remove RST files that were auto-generated for scripts and copy .py into built _sources. |
| 105 | +
|
| 106 | + Safety rules: |
| 107 | + - Only remove files whose first line starts with the auto-generated marker |
| 108 | + ('.. AUTO-GENERATED from ...'). This avoids deleting user-authored RST. |
| 109 | + - Only remove files after a successful build (exception is None). |
| 110 | + """ |
| 111 | + logger = logging.getLogger(__name__) |
| 112 | + |
| 113 | + # If build failed, do not remove files so devs can inspect outputs |
| 114 | + if exception is not None: |
| 115 | + logger.info('Build failed; leaving auto-generated RST files in place') |
| 116 | + return |
| 117 | + |
| 118 | + srcdir = app.srcdir |
| 119 | + patterns = [ |
| 120 | + os.path.join(srcdir, 'usage', 'examples', '*.rst'), |
| 121 | + ] |
| 122 | + |
| 123 | + rst_files = [] |
| 124 | + for pat in patterns: |
| 125 | + rst_files.extend(glob.glob(pat)) |
| 126 | + |
| 127 | + removed = 0 |
| 128 | + for rst in sorted(set(rst_files)): |
| 129 | + try: |
| 130 | + with io.open(rst, 'r', encoding='utf-8') as fh: |
| 131 | + first = fh.readline().strip() |
| 132 | + if first.startswith(MARKER_TEMPLATE.format('')): |
| 133 | + # If building HTML, copy the original .py into the built _sources |
| 134 | + # directory so the "View source" link can serve the real Python file. |
| 135 | + try: |
| 136 | + py_full = os.path.splitext(rst)[0] + '.py' |
| 137 | + if os.path.exists(py_full) and getattr(app, 'builder', None) is not None: |
| 138 | + try: |
| 139 | + builder_name = getattr(app.builder, 'name', None) |
| 140 | + outdir = getattr(app.builder, 'outdir', None) |
| 141 | + except Exception: |
| 142 | + builder_name = None |
| 143 | + outdir = None |
| 144 | + |
| 145 | + if builder_name == 'html' and outdir: |
| 146 | + dest_dir = os.path.join(outdir, '_sources') |
| 147 | + rel_py = os.path.relpath(py_full, app.srcdir) |
| 148 | + dest_path = os.path.join(dest_dir, rel_py) |
| 149 | + os.makedirs(os.path.dirname(dest_path), exist_ok=True) |
| 150 | + # Copy file contents |
| 151 | + try: |
| 152 | + with io.open(py_full, 'r', encoding='utf-8') as rf, io.open(dest_path, 'w', encoding='utf-8') as wf: |
| 153 | + wf.write(rf.read()) |
| 154 | + logger.info(f'Copied Python source to built _sources: {dest_path}') |
| 155 | + except Exception as e: |
| 156 | + logger.warning(f'Failed to write built source {dest_path}: {e}') |
| 157 | + |
| 158 | + except Exception: |
| 159 | + # non-fatal; continue to removal attempt |
| 160 | + pass |
| 161 | + |
| 162 | + try: |
| 163 | + os.remove(rst) |
| 164 | + removed += 1 |
| 165 | + logger.info(f'Removed generated RST: {rst}') |
| 166 | + except Exception as e: |
| 167 | + logger.warning(f'Failed to remove generated RST {rst}: {e}') |
| 168 | + except Exception as e: |
| 169 | + logger.warning(f'Could not read {rst}; skipping removal: {e}') |
| 170 | + |
| 171 | + logger.info(f'Removed {removed} auto-generated script RST files') |
| 172 | + |
| 173 | + |
| 174 | +def _rewrite_sourcelink_to_py(app, pagename, templatename, context, doctree): |
| 175 | + """If the current page was generated from a script RST, point "View source" to the .py. |
| 176 | +
|
| 177 | + This mirrors nbsphinx behaviour for notebooks where the "Show source" link |
| 178 | + points to the original notebook instead of the generated RST. |
| 179 | + """ |
| 180 | + try: |
| 181 | + # doc2path with base=None returns a path relative to srcdir |
| 182 | + rst_rel = app.env.doc2path(pagename, base=None) |
| 183 | + except Exception: |
| 184 | + return |
| 185 | + |
| 186 | + rst_full = os.path.join(app.srcdir, rst_rel) |
| 187 | + if not os.path.exists(rst_full): |
| 188 | + return |
| 189 | + |
| 190 | + try: |
| 191 | + with io.open(rst_full, 'r', encoding='utf-8') as fh: |
| 192 | + first = fh.readline().strip() |
| 193 | + except Exception: |
| 194 | + return |
| 195 | + |
| 196 | + # Only rewrite for our auto-generated script RST files |
| 197 | + if not first.startswith(MARKER_TEMPLATE.format('')): |
| 198 | + return |
| 199 | + |
| 200 | + py_full = os.path.splitext(rst_full)[0] + '.py' |
| 201 | + if not os.path.exists(py_full): |
| 202 | + return |
| 203 | + |
| 204 | + # Make the path relative to the source directory (what Sphinx expects) |
| 205 | + rel_py = os.path.relpath(py_full, app.srcdir) |
| 206 | + |
| 207 | + # Replace the template context variable so the HTML theme will link to .py |
| 208 | + context['sourcename'] = rel_py |
| 209 | + |
| 210 | + |
| 211 | +def setup(app): |
| 212 | + # Generate per-script RST files before reading sources |
| 213 | + app.connect('builder-inited', _generate_rst_for_scripts) |
| 214 | + # Remove the generated per-script RST files after a successful build |
| 215 | + app.connect('build-finished', _remove_generated_rst) |
| 216 | + # Make "View source" point to the original .py for auto-generated script pages |
| 217 | + app.connect('html-page-context', _rewrite_sourcelink_to_py) |
| 218 | + |
| 219 | + return { |
| 220 | + 'version': '1.0', |
| 221 | + 'parallel_read_safe': True, |
| 222 | + 'parallel_write_safe': True, |
| 223 | + } |
0 commit comments