Get thumbnails working for gallery (#340)

* Initial version of gallery using nbsphinx with a hack to make it work with myst_nb.

* Add gallery.

* Run black.

* Add all dirs to gallery.

* WIP: Trying to get thumbnails to work.

* Remove sphinxext stuff

* Remove print and debug stuff.

* Remove unused method.

* fix readthedocs config

* sort inputs and requirements

Co-authored-by: Oriol (ZBook) <[email protected]>

Author: twiecki

.readthedocs.yml

@@ -7,5 +7,3 @@ python:
    version: "3.8"
    install:
    - requirements: requirements-docs.txt
-   - method: pip
-     path: .

examples/conf.py

@@ -1,4 +1,4 @@
-import os
+import os, sys
 from sphinx.application import Sphinx
 
 # -- Project information -----------------------------------------------------
@@ -8,6 +8,8 @@
 
 # -- General configuration ---------------------------------------------------
 
+sys.path.insert(0, os.path.abspath("../sphinxext"))
+
 # Add any Sphinx extension module names here, as strings. They can be
 # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
 # ones.
@@ -23,6 +25,7 @@
     "sphinx_codeautolink",
     "notfound.extension",
     "sphinx_gallery.load_style",
+    "gallery_generator",
 ]
 
 # List of patterns, relative to source directory, that match files and
@@ -46,22 +49,19 @@ def hack_nbsphinx(app: Sphinx) -> None:
         GalleryNode,
         NbGallery,
         patched_toctree_resolve,
-        NotebookParser,
-        NbInput,
-        NbOutput,
-        NbInfo,
-        NbWarning,
-        CodeAreaNode,
-        depart_codearea_html,
-        visit_codearea_latex,
-        depart_codearea_latex,
-        GetSizeFromImages,
     )
     from sphinx.environment.adapters import toctree
 
-    nbsphinx_thumbnails = {
-        "case_studies/stochastic_volatility": "_static/stochastic_volatility.png",
-    }
+    from glob import glob
+
+    nb_paths = glob("examples/*/*.ipynb")
+    nbsphinx_thumbnails = {}
+    for nb_path in nb_paths:
+        png_file = os.path.join("_static", os.path.splitext(os.path.split(nb_path)[-1])[0] + ".png")
+        nb_path_rel = os.path.splitext(
+            os.path.join(*os.path.normpath(nb_path).split(os.path.sep)[1:])
+        )[0]
+        nbsphinx_thumbnails[nb_path_rel] = png_file
 
     def builder_inited(app: Sphinx):
         if not hasattr(app.env, "nbsphinx_thumbnails"):
@@ -70,46 +70,6 @@ def builder_inited(app: Sphinx):
     def do_nothing(*node):
         pass
 
-    app.add_source_parser(NotebookParser)
-    # app.add_config_value('nbsphinx_execute', 'auto', rebuild='env')
-    app.add_config_value("nbsphinx_kernel_name", "", rebuild="env")
-    app.add_config_value("nbsphinx_execute_arguments", [], rebuild="env")
-    app.add_config_value("nbsphinx_allow_errors", False, rebuild="")
-    app.add_config_value("nbsphinx_timeout", None, rebuild="")
-    app.add_config_value("nbsphinx_codecell_lexer", "none", rebuild="env")
-    app.add_config_value("nbsphinx_prompt_width", "4.5ex", rebuild="html")
-    app.add_config_value("nbsphinx_responsive_width", "540px", rebuild="html")
-    app.add_config_value("nbsphinx_prolog", None, rebuild="env")
-    app.add_config_value("nbsphinx_epilog", None, rebuild="env")
-    app.add_config_value("nbsphinx_input_prompt", "[%s]:", rebuild="env")
-    app.add_config_value("nbsphinx_output_prompt", "[%s]:", rebuild="env")
-    app.add_config_value("nbsphinx_custom_formats", {}, rebuild="env")
-    # Default value is set in config_inited():
-    app.add_config_value("nbsphinx_requirejs_path", None, rebuild="html")
-    # Default value is set in config_inited():
-    app.add_config_value("nbsphinx_requirejs_options", None, rebuild="html")
-    # This will be updated in env_updated():
-    app.add_config_value("nbsphinx_widgets_path", None, rebuild="html")
-    app.add_config_value("nbsphinx_widgets_options", {}, rebuild="html")
-    # app.add_config_value('nbsphinx_thumbnails', {}, rebuild='html')
-    app.add_config_value("nbsphinx_assume_equations", True, rebuild="env")
-
-    app.add_directive("nbinput", NbInput)
-    app.add_directive("nboutput", NbOutput)
-    app.add_directive("nbinfo", NbInfo)
-    app.add_directive("nbwarning", NbWarning)
-    app.add_directive("nbgallery", NbGallery)
-    app.add_node(
-        CodeAreaNode,
-        html=(do_nothing, depart_codearea_html),
-        latex=(visit_codearea_latex, depart_codearea_latex),
-        text=(do_nothing, do_nothing),
-    )
-    app.connect("builder-inited", builder_inited)
-    app.connect("doctree-resolved", doctree_resolved)
-    app.add_post_transform(GetSizeFromImages)
-
-    app.add_config_value("nbsphinx_execute", "auto", rebuild="env")
     app.add_config_value("nbsphinx_thumbnails", nbsphinx_thumbnails, rebuild="html")
     app.add_directive("nbgallery", NbGallery)
     app.add_node(
@@ -185,7 +145,7 @@ def setup(app: Sphinx):
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
 # so a file named "default.css" will overwrite the builtin "default.css".
-html_static_path = ["../_static", "../thumbnails"]
+html_static_path = ["../_static", "../_images", "../_templates"]
 html_css_files = ["custom.css"]
 templates_path = ["../_templates"]
 html_sidebars = {

examples/gallery.rst

@@ -112,4 +112,4 @@ Variational Inference
     :glob:
     :reversed:
 
-    variational_inference/*
+    variational_inference/*

requirements-docs.txt

@@ -10,3 +10,4 @@ sphinx-codeautolink
 sphinx-notfound-page
 nbsphinx
 sphinx-gallery
+matplotlib

sphinxext/gallery_generator.py

@@ -0,0 +1,94 @@
+"""
+Sphinx plugin to run generate a gallery for notebooks
+
+Modified from the seaborn project, which modified the mpld3 project.
+"""
+import base64
+import json
+import os
+import shutil
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+from matplotlib import image
+
+DOC_SRC = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+DEFAULT_IMG_LOC = os.path.join(os.path.dirname(DOC_SRC), "pymc-examples/_static", "PyMC.png")
+
+
+def create_thumbnail(infile, width=275, height=275, cx=0.5, cy=0.5, border=4):
+    """Overwrites `infile` with a new file of the given size"""
+    im = image.imread(infile)
+    rows, cols = im.shape[:2]
+    size = min(rows, cols)
+    if size == cols:
+        xslice = slice(0, size)
+        ymin = min(max(0, int(cx * rows - size // 2)), rows - size)
+        yslice = slice(ymin, ymin + size)
+    else:
+        yslice = slice(0, size)
+        xmin = min(max(0, int(cx * cols - size // 2)), cols - size)
+        xslice = slice(xmin, xmin + size)
+    thumb = im[yslice, xslice]
+    thumb[:border, :, :3] = thumb[-border:, :, :3] = 0
+    thumb[:, :border, :3] = thumb[:, -border:, :3] = 0
+
+    dpi = 100
+    fig = plt.figure(figsize=(width / dpi, height / dpi), dpi=dpi)
+
+    ax = fig.add_axes([0, 0, 1, 1], aspect="auto", frameon=False, xticks=[], yticks=[])
+    ax.imshow(thumb, aspect="auto", resample=True, interpolation="bilinear")
+    fig.savefig(infile, dpi=dpi)
+    plt.close(fig)
+    return fig
+
+
+class NotebookGenerator:
+    """Tools for generating an example page from a file"""
+
+    def __init__(self, filename, target_dir):
+        self.basename = os.path.basename(filename)
+        stripped_name = os.path.splitext(self.basename)[0]
+        self.image_dir = os.path.join(target_dir, "_static")
+        self.png_path = os.path.join(self.image_dir, f"{stripped_name}.png")
+        with open(filename) as fid:
+            self.json_source = json.load(fid)
+        self.default_image_loc = DEFAULT_IMG_LOC
+
+    def extract_preview_pic(self):
+        """By default, just uses the last image in the notebook."""
+        pic = None
+        for cell in self.json_source["cells"]:
+            for output in cell.get("outputs", []):
+                if "image/png" in output.get("data", []):
+                    pic = output["data"]["image/png"]
+        if pic is not None:
+            return base64.b64decode(pic)
+        return None
+
+    def gen_previews(self):
+        preview = self.extract_preview_pic()
+        if preview is not None:
+            with open(self.png_path, "wb") as buff:
+                buff.write(preview)
+        else:
+            print(f"didn't find for {self.png_path}")
+            shutil.copy(self.default_image_loc, self.png_path)
+        create_thumbnail(self.png_path)
+
+
+def main(app):
+    from glob import glob
+
+    nb_paths = glob("examples/*/*.ipynb")
+
+    for nb_path in nb_paths:
+        nbg = NotebookGenerator(nb_path, "")
+        nbg.gen_previews()
+
+
+def setup(app):
+    app.connect("builder-inited", main)