datasette-seaborn/datasette_seaborn/__init__.py

from datasette import hookimpl
from datasette.utils.asgi import Response
import seaborn as sns
import pathlib
import asyncio, sys
import json
 
render_code = (pathlib.Path(__file__).parent / "render.py").read_text()
 
DEFAULT_TIME_LIMIT = 5.0
 
 
def error(request, msg):
    return {"body": msg, "status_code": 500}
 
 
async def render_seaborn(datasette, columns, rows, request):
    # _seaborn= provides the method
    # _seaborn_X provide the arguments
    method = request.args.get("_seaborn") or "NOMETHOD"
    if not hasattr(sns, method):
        return error(request, "_seaborn= is required")
 
    kwargs = {
        key.replace("_seaborn_", ""): request.args[key]
        for key in request.args
        if key.startswith("_seaborn_")
    }
 
    plugin_config = datasette.plugin_config("datasette-seaborn") or {}
    render_time_limit = plugin_config.get("render_time_limit") or DEFAULT_TIME_LIMIT
 
    proc = await asyncio.create_subprocess_exec(
        sys.executable,
        "-",
        stdout=asyncio.subprocess.PIPE,
        stdin=asyncio.subprocess.PIPE,
    )
 
    code = (
        render_code
        + "\n\n"
        + "null=None\n"
        + "true=True\n"
        + "false=False\n"
        + "render_records({records}, {method}, {kwargs})".format(
            records=json.dumps([dict(row) for row in rows]),
            method=json.dumps(method),
            kwargs=json.dumps(kwargs),
        )
    )
 
    image_bytes = []
    errors = []
 
    async def render_image():
        stdout_data, stderr_data = await proc.communicate(code.encode("utf-8"))
        if stdout_data:
            image_bytes.append(stdout_data)
        if stderr_data:
            errors.append(stdout_data)
 
    time_limit_hit = False
    # Run this with a time limit
    try:
        await asyncio.wait_for(render_image(), timeout=render_time_limit)
    except asyncio.TimeoutError:
        time_limit_hit = True
    try:
        proc.kill()
    except OSError:
        # Ignore 'no such process' error
        pass
 
    if time_limit_hit:
        return Response("Render took too long", status=500)
    elif errors:
        return Response("Errors occurred: {}".format(errors[0]), status=500)
    else:
        return Response(image_bytes[0], content_type="image/png")
 
 
@hookimpl
def register_output_renderer():
    return {
        "extension": "seaborn",
        "render": render_seaborn,
    }