Merge pull request #3 from compute-tooling/add-screenshot

Add module for taking screenshots of renderable outputs

MANIFEST.in

@@ -0,0 +1 @@
+include cs_storage/templates/index.html

cs_storage/__init__.py

@@ -11,6 +11,8 @@
 from marshmallow import Schema, fields, validate
 
 
+from .screenshot import screenshot, ScreenshotError, SCREENSHOT_ENABLED
+
 __version__ = "1.7.0"
 
 
@@ -21,6 +23,7 @@ class Serializer:
     """
     Base class for serializng input data to bytes and back.
     """
+
     def __init__(self, ext):
         self.ext = ext
 
@@ -77,9 +80,24 @@ def get_serializer(media_type):
 class Output:
     """Output mixin shared among LocalOutput and RemoteOutput"""
 
+    id = fields.UUID(required=False)
     title = fields.Str()
     media_type = fields.Str(
-        validate=validate.OneOf(choices=["bokeh", "table", "CSV", "PNG", "JPEG", "MP3", "MP4", "HDF5", "PDF", "Markdown", "Text"])
+        validate=validate.OneOf(
+            choices=[
+                "bokeh",
+                "table",
+                "CSV",
+                "PNG",
+                "JPEG",
+                "MP3",
+                "MP4",
+                "HDF5",
+                "PDF",
+                "Markdown",
+                "Text",
+            ]
+        )
     )
 
 
@@ -111,6 +129,28 @@ class LocalResult(Schema):
     downloadable = fields.Nested(LocalOutput, many=True)
 
 
+def write_pic(fs, output):
+    if SCREENSHOT_ENABLED:
+        s = time.time()
+        try:
+            pic_data = screenshot(output)
+        except ScreenshotError:
+            print("failed to create screenshot for ", output["id"])
+            return
+        else:
+            with fs.open(f"{BUCKET}/{output['id']}.png", "wb") as f:
+                f.write(pic_data)
+            f = time.time()
+            print(f"Pic write finished in {f-s}s")
+    else:
+        import warnings
+
+        warnings.warn(
+            "Screenshot not enabled. Make sure you have installed "
+            "the optional packages listed in environment.yaml."
+        )
+
+
 def write(task_id, loc_result, do_upload=True):
     fs = gcsfs.GCSFileSystem()
     s = time.time()
@@ -124,17 +164,21 @@ def write(task_id, loc_result, do_upload=True):
         for output in loc_result[category]:
             serializer = get_serializer(output["media_type"])
             ser = serializer.serialize(output["data"])
+            output["id"] = str(uuid.uuid4())
             filename = output["title"]
             if not filename.endswith(f".{serializer.ext}"):
                 filename += f".{serializer.ext}"
             zipfileobj.writestr(filename, ser)
             rem_result[category]["outputs"].append(
                 {
+                    "id": output["id"],
                     "title": output["title"],
                     "media_type": output["media_type"],
                     "filename": filename,
                 }
             )
+            if do_upload and category == "renderable":
+                write_pic(fs, output)
         zipfileobj.close()
         buff.seek(0)
         if do_upload:
@@ -160,9 +204,12 @@ def read(rem_result, json_serializable=True):
 
         for rem_output in rem_result[category]["outputs"]:
             ser = get_serializer(rem_output["media_type"])
-            rem_data = ser.deserialize(zipfileobj.read(rem_output["filename"]), json_serializable)
+            rem_data = ser.deserialize(
+                zipfileobj.read(rem_output["filename"]), json_serializable
+            )
             read[category].append(
                 {
+                    "id": rem_output.get("id", None),
                     "title": rem_output["title"],
                     "media_type": rem_output["media_type"],
                     "data": rem_data,
@@ -171,3 +218,11 @@ def read(rem_result, json_serializable=True):
     f = time.time()
     print(f"Read finished in {f-s}s")
     return read
+
+
+def add_screenshot_links(rem_result):
+    for rem_output in rem_result["renderable"]["outputs"]:
+        rem_output[
+            "screenshot"
+        ] = f"https://storage.googleapis.com/{BUCKET}/{rem_output['id']}.png"
+    return rem_result

cs_storage/screenshot.py

@@ -0,0 +1,111 @@
+import asyncio
+import os
+import tempfile
+
+try:
+    # These dependencies are optional. The storage component may be used
+    # without the screenshot component.
+    from jinja2 import Template
+    from pyppeteer import launch
+    from bokeh.resources import CDN
+
+    BASE_ARGS = {
+        "bokeh_scripts": {"cdn_js": CDN.js_files[0], "widget_js": CDN.js_files[1]}
+    }
+    SCREENSHOT_ENABLED = True
+
+except ImportError:
+    SCREENSHOT_ENABLED = False
+    Template = None
+    launch = None
+    CDN = None
+    BASE_ARGS = {}
+
+import cs_storage
+
+
+CURRENT_DIR = os.path.abspath(os.path.dirname(__file__))
+
+
+class ScreenshotError(Exception):
+    pass
+
+
+def get_template():
+    if not SCREENSHOT_ENABLED:
+        return None
+    with open(f"{CURRENT_DIR}/templates/index.html", "r") as f:
+        text = f.read()
+
+    template = Template(text)
+
+    return template
+
+
+TEMPLATE = get_template()
+
+
+def write_template(output):
+    kwargs = {**BASE_ARGS, **{"output": output}}
+    return TEMPLATE.render(**kwargs)
+
+
+async def _screenshot(template_path, pic_path):
+    """
+    Use pyppeteer, a python port of puppeteer, to open the
+    template at template_path and take a screenshot of the
+    output that is rendered within it.
+
+    The output is rendered within a Bootstrap card element.
+    This element is only as big as the elements that it contains.
+    Thus, we only need to get the dimensions of the bootstrap
+    card to figure out which part of the screen we need to use
+    for the screenshot!
+
+    Note: pyppetter looks stale. If it continues to not be
+    maintained well, then the extremely active, well-maintained
+    puppeteer should be used for creating these screenshots. The
+    downside of using puppeteer is that it is written in nodejs.
+    """
+    browser = await launch(args=["--no-sandbox"])
+    page = await browser.newPage()
+    await page.goto(f"file://{template_path}")
+    await page.setViewport(dict(width=1920, height=1080))
+    await page.waitFor(1000)
+    element = await page.querySelector("#output")
+    if element is None:
+        raise ScreenshotError("Unable to take screenshot.")
+    boundingbox = await element.boundingBox()
+    clip = dict(
+        x=boundingbox["x"],
+        y=boundingbox["y"],
+        width=min(boundingbox["width"], 1920),
+        height=min(boundingbox["height"], 1080),
+    )
+    await page.screenshot(path=f"{pic_path}", type="png", clip=clip)
+    await browser.close()
+
+
+def screenshot(output, debug=False):
+    """
+    Create screenshot of outputs. The intermediate results are
+    written to temporary files and a picture, represented as a
+    stream of bytes, is returned.
+    """
+    if not SCREENSHOT_ENABLED:
+        return None
+    html = write_template(output)
+    with tempfile.NamedTemporaryFile(suffix=".html") as temp:
+        if debug:
+            with open(f'{output["title"]}.html', "w") as f:
+                f.write(html)
+        temp.write(html.encode("utf-8"))
+        temp.seek(0)
+        template_path = temp.name
+        with tempfile.NamedTemporaryFile(suffix=".png") as pic:
+            pic_path = pic.name
+            asyncio.get_event_loop().run_until_complete(
+                _screenshot(template_path, pic_path)
+            )
+            pic_bytes = pic.read()
+    return pic_bytes

cs_storage/templates/index.html

@@ -0,0 +1,60 @@
+<!DOCTYPE html>
+<html lang="en">
+
+<head>
+  <title>Compute Studio</title>
+  <script src="https://code.jquery.com/jquery-3.3.1.min.js"
+    integrity="sha256-FgpCb/KJQlLNfOu91ta32o/NMZxltwRo8QtmkMRdAu8=" crossorigin="anonymous"></script>
+  <script src="https://cdnjs.cloudflare.com/ajax/libs/popper.js/1.14.3/umd/popper.min.js"
+    integrity="sha384-ZMP7rVo3mIykV+2+9J3UJ46jBk0WLaUAdn689aCwoqbBJiSnjAK/l8WvCWPIPm49"
+    crossorigin="anonymous"></script>
+  <link rel="stylesheet" href="https://stackpath.bootstrapcdn.com/bootstrap/4.1.3/css/bootstrap.min.css"
+    integrity="sha384-MCw98/SFnGE8fJT3GXwEOngsV7Zt27NXFoaoApmYm81iuXoPkFOJwJ8ERdknLPMO" crossorigin="anonymous">
+  <script src="https://stackpath.bootstrapcdn.com/bootstrap/4.1.3/js/bootstrap.min.js"
+    integrity="sha384-ChfqqxuZUCnJSK3+MXmPNIyE6ZbWh2IMqE241rYiqJxyMiZ6OW/JmZQ5stwEULTy"
+    crossorigin="anonymous"></script>
+  <!-- <link rel="stylesheet" href="https://use.fontawesome.com/releases/v5.5.0/css/all.css" integrity="sha384-B4dIYHKNBt8Bc12p+WXckhzcICo0wtJAoU8YZTY5qE0Id1GSseTk6S+L3BlXeVIU" crossorigin="anonymous"> -->
+  <link href="https://fonts.googleapis.com/css?family=Roboto" rel="stylesheet">
+
+  <meta name="viewport" content="width=device-width, initial-scale=1">
+
+
+  <script type="text/javascript" src={{bokeh_scripts.cdn_js|safe}}></script>
+  <script type="text/javascript" src={{bokeh_scripts.widget_js|safe}}></script>
+
+</head>
+
+<body>
+  <div class="container mt-4">
+    <div class="col-md-12">
+      <div id="output" class="card card-body card-inner text-center" style="overflow:auto">
+        <h4>{{output.title}}</h4>
+        {% if output.media_type == 'bokeh' %}
+        <div style="margin:0 auto;">{{output.id}}
+          <div id="{{output.id}}" data-root-id="" className="bk-root"></div>
+        </div>
+        {% elif output.media_type == 'table' %}
+        <div style="margin:0 auto;">
+          {{output.data|safe}}
+        </div>
+        {% elif output.media_type == 'PNG' %}
+        <div style="margin:0 auto;">
+          <img src="data:image/png;base64, {{ output.data|safe }}" />
+        </div>
+        {% elif output.media_type == 'JPEG' %}
+        <div style="margin:0 auto;">
+          <img src="data:image/jpeg;base64, {{ output.data|safe }}" />
+        </div>
+        {% endif %}
+      </div>
+    </div>
+  </div>
+</body>
+
+{% if output.media_type == "bokeh" %}
+<script>
+  window.Bokeh.embed.embed_item({{ output.data | tojson | safe }}, "{{ output.id }}");
+</script>
+{% endif %}
+
+</html>