image

packages/python/yap_kernel/images.py

@@ -0,0 +1,48 @@
+import base64
+import imghdr
+import os
+
+#from IPython.
+
+_TEXT_SAVED_IMAGE = "yap_kernel: saved image data to:"
+
+image_setup_cmd = """
+display () {
+    TMPFILE=$(mktemp ${TMPDIR-/tmp}/yap_kernel.XXXXXXXXXX)
+    cat > $TMPFILE
+    echo "%s $TMPFILE" >&2
+}
+""" % _TEXT_SAVED_IMAGE
+
+def display_data_for_image(filename):
+    with open(filename, 'rb') as f:
+        image = f.read()
+    os.unlink(filename)
+
+    image_type = imghdr.what(None, image)
+    if image_type is None:
+        raise ValueError("Not a valid image: %s" % image)
+
+    image_data = base64.b64encode(image).decode('ascii')
+    content = {
+        'data': {
+            'image/' + image_type: image_data
+        },
+        'metadata': {}
+    }
+    return content
+
+
+def extract_image_filenames(output):
+    output_lines = []
+    image_filenames = []
+
+    for line in output.split("\n"):
+        if line.startswith(_TEXT_SAVED_IMAGE):
+            filename = line.rstrip().split(": ")[-1]
+            image_filenames.append(filename)
+        else:
+            output_lines.append(line)
+
+    output = "\n".join(output_lines)
+    return image_filenames, output