1""" 2Sphinx plugin to run generate a gallery for notebooks 3 4Modified from the seaborn project, which modified the mpld3 project. 5""" 6import base64 7import json 8import os 9import runpy 10import shutil 11 12from pathlib import Path 13 14import matplotlib 15 16matplotlib.use("Agg") 17import matplotlib.pyplot as plt 18 19from matplotlib import image 20 21DOC_SRC = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) 22DEFAULT_IMG_LOC = os.path.join(os.path.dirname(DOC_SRC), "logos", "PyMC3.png") 23TABLE_OF_CONTENTS_FILENAME = "table_of_contents_{}.js" 24 25INDEX_TEMPLATE = """ 26:orphan: 27 28.. 29 _href from docs/source/conf.py 30 31.. _{sphinx_tag}: 32 33.. title:: {gallery}_notebooks 34 35.. raw:: html 36 37 <h1 class="ui header">{Gallery} Notebooks</h1> 38 <div id="gallery" class="ui vertical segment"> 39 </div> 40""" 41 42 43def create_thumbnail(infile, width=275, height=275, cx=0.5, cy=0.5, border=4): 44 """Overwrites `infile` with a new file of the given size""" 45 im = image.imread(infile) 46 rows, cols = im.shape[:2] 47 size = min(rows, cols) 48 if size == cols: 49 xslice = slice(0, size) 50 ymin = min(max(0, int(cx * rows - size // 2)), rows - size) 51 yslice = slice(ymin, ymin + size) 52 else: 53 yslice = slice(0, size) 54 xmin = min(max(0, int(cx * cols - size // 2)), cols - size) 55 xslice = slice(xmin, xmin + size) 56 thumb = im[yslice, xslice] 57 thumb[:border, :, :3] = thumb[-border:, :, :3] = 0 58 thumb[:, :border, :3] = thumb[:, -border:, :3] = 0 59 60 dpi = 100 61 fig = plt.figure(figsize=(width / dpi, height / dpi), dpi=dpi) 62 63 ax = fig.add_axes([0, 0, 1, 1], aspect="auto", frameon=False, xticks=[], yticks=[]) 64 ax.imshow(thumb, aspect="auto", resample=True, interpolation="bilinear") 65 fig.savefig(infile, dpi=dpi) 66 plt.close(fig) 67 return fig 68 69 70class NotebookGenerator: 71 """Tools for generating an example page from a file""" 72 73 def __init__(self, filename, target_dir): 74 self.basename = os.path.basename(filename) 75 stripped_name = os.path.splitext(self.basename)[0] 76 self.output_html = str(".." / Path(filename).relative_to(Path.cwd()).with_suffix(".html")) 77 self.image_dir = os.path.join(target_dir, "_images") 78 self.png_path = os.path.join(self.image_dir, f"{stripped_name}.png") 79 with open(filename) as fid: 80 self.json_source = json.load(fid) 81 self.pagetitle = self.extract_title() 82 self.default_image_loc = DEFAULT_IMG_LOC 83 84 # Only actually run it if the output RST file doesn't 85 # exist or it was modified less recently than the example 86 if not os.path.exists(self.output_html) or ( 87 os.path.getmtime(self.output_html) < os.path.getmtime(filename) 88 ): 89 90 self.gen_previews() 91 else: 92 print(f"skipping {filename}") 93 94 def extract_preview_pic(self): 95 """By default, just uses the last image in the notebook.""" 96 pic = None 97 for cell in self.json_source["cells"]: 98 for output in cell.get("outputs", []): 99 if "image/png" in output.get("data", []): 100 pic = output["data"]["image/png"] 101 if pic is not None: 102 return base64.b64decode(pic) 103 return None 104 105 def extract_title(self): 106 for cell in self.json_source["cells"]: 107 if cell["cell_type"] == "markdown": 108 rows = [row.strip() for row in cell["source"] if row.strip()] 109 for row in rows: 110 if row.startswith("# "): 111 return row[2:] 112 return self.basename.replace("_", " ") 113 114 def gen_previews(self): 115 preview = self.extract_preview_pic() 116 if preview is not None: 117 with open(self.png_path, "wb") as buff: 118 buff.write(preview) 119 else: 120 shutil.copy(self.default_image_loc, self.png_path) 121 create_thumbnail(self.png_path) 122 123 124class TableOfContentsJS: 125 """Container to load table of contents JS file""" 126 127 def load(self, path): 128 """Creates an attribute ``contents`` by running the JS file as a python 129 file. 130 131 """ 132 runpy.run_path(path, {"Gallery": self}) 133 134 135def build_gallery(srcdir, gallery): 136 working_dir = os.getcwd() 137 os.chdir(srcdir) 138 static_dir = os.path.join(srcdir, "_static") 139 target_dir = os.path.join(srcdir, f"nb_{gallery}") 140 image_dir = os.path.join(target_dir, "_images") 141 source_dir = os.path.abspath( 142 os.path.join( 143 os.path.dirname(os.path.dirname(srcdir)), "docs", "source", "pymc-examples", "examples" 144 ) 145 ) 146 table_of_contents_file = os.path.join(source_dir, TABLE_OF_CONTENTS_FILENAME.format(gallery)) 147 tocjs = TableOfContentsJS() 148 tocjs.load(table_of_contents_file) 149 150 if not os.path.exists(static_dir): 151 os.makedirs(static_dir) 152 153 if not os.path.exists(target_dir): 154 os.makedirs(target_dir) 155 156 if not os.path.exists(image_dir): 157 os.makedirs(image_dir) 158 159 if not os.path.exists(source_dir): 160 os.makedirs(source_dir) 161 162 # Create default image 163 default_png_path = os.path.join(os.path.join(target_dir, "_images"), "default.png") 164 shutil.copy(DEFAULT_IMG_LOC, default_png_path) 165 create_thumbnail(default_png_path) 166 167 # Write individual example files 168 data = {} 169 for basename in sorted(tocjs.contents): 170 if basename.find(".rst") < 1: 171 filename = os.path.join(source_dir, basename + ".ipynb") 172 ex = NotebookGenerator(filename, target_dir) 173 url = Path(os.path.join(os.sep, gallery, ex.output_html)) 174 # Need to chop off "/${gallery}/../" so as redirection works in multi versioned docs. 175 url = str(Path("..", *url.parts[3:])) 176 data[basename] = { 177 "title": ex.pagetitle, 178 "url": url, 179 "thumb": os.path.basename(ex.png_path), 180 } 181 182 else: 183 filename = basename.split(".")[0] 184 url = Path(os.path.join(os.sep, gallery, "../" + filename + ".html")) 185 # Need to chop off "/${gallery}/../" so as redirection works in multi versioned docs. 186 url = str(Path("..", *url.parts[3:])) 187 data[basename] = { 188 "title": " ".join(filename.split("_")), 189 "url": url, 190 "thumb": os.path.basename(default_png_path), 191 } 192 193 js_file = os.path.join(image_dir, f"gallery_{gallery}_contents.js") 194 with open(table_of_contents_file) as toc: 195 table_of_contents = toc.read() 196 197 js_contents = "Gallery.examples = {}\n{}".format(json.dumps(data), table_of_contents) 198 199 with open(js_file, "w") as js: 200 js.write(js_contents) 201 202 with open(os.path.join(target_dir, "index.rst"), "w") as index: 203 index.write( 204 INDEX_TEMPLATE.format( 205 sphinx_tag="notebook_gallery", gallery=gallery, Gallery=gallery.title().rstrip("s") 206 ) 207 ) 208 209 os.chdir(working_dir) 210 211 212def main(app): 213 for gallery in ("tutorials", "examples"): 214 build_gallery(app.builder.srcdir, gallery) 215 216 217def setup(app): 218 app.connect("builder-inited", main) 219