render.py

#!/usr/local/bin/python
# encoding=utf8

import sys
import json
from jinja2 import Environment, FileSystemLoader
from collections import defaultdict
import re
import markdown
import logging
import subprocess
import base64
import hashlib
import io
from PIL import Image


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
md = markdown.Markdown(extensions=['meta', 'footnotes'])

packages = {
    "pandas": "Pandas",
    "matplotlib": "Matplotlib",
    "seaborn": "Seaborn",
    "seaborn-objects": "seaborn.objects",
    "plotnine": "plotnine",
    "lets-plot": "lets-plot",
    "plotly": "plotly",
    "hvplot": "hvPlot (Bokeh)",
    "altair": "Altair",
    "ggplot": "ggplot2 (R)",
}

names = {
    "bar-counts": "Bar Chart",
    "simple-histogram": "Histogram",
    "box-plot": "Box Plot",
    "scatter-plot": "Scatter Plot",
    "timeseries": "Time Series",
    "scatter-plot-with-colors": "Scatter Plot with Faceted with Color",
    "scatter-plot-with-size": "Scatter Plot with Points Sized by Continuous Value",
    "scatter-plot-with-facet": "Scatter Plot Faceted on One Variable",
    "scatter-plot-with-facets": "Scatter Plot Faceted on Two Variables",
    "scatter-with-regression": "Scatter Plot and Regression Line with 95% Confidence Interval Layered",
    "stacked-smooth-line-and-scatter": "Smoothed Line Plot and Scatter Plot Layered",
    "stacked-bar-chart": "Stacked Bar Chart",
    "dodged-bar-chart": "Dodged Bar Chart",
    "stacked-kde": "Stacked KDE Plot",
    "heatmap": "Heatmap",
}

with open("INTRO.md", "r") as f:
    intro = f.read()


WEBP_QUALITY = 85


def image_from_cell(cell):
    """Save the cell's PNG output as WebP and return its path and pixel size.

    The filename stays the MD5 of the base64 PNG, so a plot that hasn't changed
    keeps its URL. The pixel size lets the template declare width/height, which
    reserves layout space for the lazily loaded images.
    """
    try:
        for c in cell['outputs']:
            if 'data' in c and 'image/png' in c['data']:
                base64_img = c['data']['image/png'].replace("\n", "").strip()
                filename = hashlib.md5()
                filename.update(base64_img.encode('ascii'))
                web_path = "/img/plots/{}.webp".format(filename.hexdigest())
                full_path = "web" + web_path
                image = Image.open(io.BytesIO(base64.b64decode(base64_img)))
                if image.mode not in ("RGB", "RGBA"):
                    image = image.convert("RGBA")
                image.save(full_path, "WEBP", quality=WEBP_QUALITY, method=6)
                return {"path": web_path, "width": image.width, "height": image.height}
    except KeyError as e:
        logging.error("Can't find image in cell: %s", cell['source'])
        raise e
    raise Exception("Can't find an image in cell %s", cell['source'])


def source_from_cell(cell):
    source = "".join(cell['source']).strip()
    source = source.replace(";", "")
    source = re.sub(r"\bImage\(.*\)", "", source) # remove bokeh render
    if "%%R" in source:
        source = '\n'.join(source.split('\n')[1:])
    if "%%altair" in source:
        source = '\n'.join(source.split('\n')[1:])
    else:
        source = source.replace('"', "'")
    if source.startswith('"""') or source.startswith("'''"):
        m = re.match("(?:[\"']{3,})((?:.|\n)*)(?:[\"']{3,})((?:.|\n)*)", source, re.MULTILINE)
        return m.groups()
    return "", source


def tags_from_cell(cell, type='ex'):
    tags = set(cell['metadata'].get('tags') or {})
    if type in tags:
        return {t.split(":")[0]: t.split(":")[1] for t in tags if ":" in t}


def data_from_cell(cell):
    classes = "table table-sm table-striped table-responsive table-bordered"
    try:
        for c in cell['outputs']:
            if 'data' in c and 'text/html' in c['data']:
                table = ' '.join(c['data']['text/html'])
                table = table.replace('border="1" class="dataframe"', 'class="{}"'.format(classes))
                table = table.replace('<thead>', '<thead class="thead-inverse">')
                return table
    except KeyError as e:
        logging.error("Can't find data in cell: %s", cell['source'])
        raise e
    raise Exception("Can't find an dataset in cell %s", cell['source'])


def reorder_meta(meta):
    def order_plots(plots):
        if plots:
            return sorted(plots, key=lambda k: list(packages.keys()).index(k['package-slug']))
        else:
            return plots
    meta = {(name, slug): order_plots(meta[slug]) for slug, name in names.items()}
    return meta


def extract_data(path):
    with open(path, 'r') as f:
        nb = json.load(f)
    cells = nb['cells']
    full_data = {}
    tags = {i: tags_from_cell(c, type='data') for i, c in enumerate(cells)}
    for cell_num, tags in tags.items():
        if tags is None:
            continue
        data = data_from_cell(cells[cell_num])
        full_data[tags['name']] = data
    return full_data


def extract_cells(path):
    with open(path, 'r') as f:
        nb = json.load(f)
    cells = nb['cells']
    tags = {i: tags_from_cell(c) for i, c in enumerate(cells)}

    meta = defaultdict(list)
    for cell_num, tags in tags.items():
        if tags is None:
            continue
        comment, source = source_from_cell(cells[cell_num])
        image = image_from_cell(cells[cell_num])

        meta[tags['name']].append({
            "cell_num": cell_num,
            "package": packages.get(tags["package"], tags["package"]),
            "package-slug": tags['package'],
            "image": image["path"],
            "image-width": image["width"],
            "image-height": image["height"],
            "content": source,
            "comment": md.convert(comment) or None,
        })
    meta = reorder_meta(meta)

    return meta


def get_git_revision_short_hash():
    return subprocess.check_output(['git', 'rev-parse', '--short', 'HEAD']).strip().decode("utf8")


if __name__ == '__main__':
    plots = extract_cells(sys.argv[1])
    data = extract_data(sys.argv[1])

    env = Environment(loader=FileSystemLoader('templates'), extensions=['jinja2_highlight.HighlightExtension'])
    template = env.get_template('t_index.html')
    output_from_parsed_template = template.render(intro=md.convert(intro),
                                                  plots=plots,
                                                  git=get_git_revision_short_hash(),
                                                  data=data)


    # to save the results
    with open("web/index.html", "w") as fh:
        fh.write(output_from_parsed_template)