diff --git a/packages/python/yap_kernel/images.py b/packages/python/yap_kernel/images.py new file mode 100644 index 000000000..f7abfbb12 --- /dev/null +++ b/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