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