diff --git a/IPython/utils/jsonutil.py b/IPython/utils/jsonutil.py index caa39fb..945f2bb 100644 --- a/IPython/utils/jsonutil.py +++ b/IPython/utils/jsonutil.py @@ -14,6 +14,7 @@ import re import sys import types +from base64 import encodestring from datetime import datetime from IPython.utils import py3compat @@ -89,6 +90,24 @@ def date_default(obj): raise TypeError("%r is not JSON serializable"%obj) +# constants for identifying png/jpeg data +PNG = b'\x89PNG\r\n\x1a\n' +JPEG = b'\xff\xd8' + +def encode_images(format_dict): + """b64-encodes images in a displaypub format dict + + Perhaps this should be handled in json_clean itself? + """ + encoded = format_dict.copy() + pngdata = format_dict.get('image/png') + if isinstance(pngdata, bytes) and pngdata[:8] == PNG: + encoded['image/png'] = encodestring(pngdata).decode('ascii') + jpegdata = format_dict.get('image/jpeg') + if isinstance(jpegdata, bytes) and jpegdata[:2] == JPEG: + encoded['image/jpeg'] = encodestring(jpegdata).decode('ascii') + return encoded + def json_clean(obj): """Clean an object to ensure it's safe to encode in JSON. diff --git a/IPython/utils/tests/test_jsonutil.py b/IPython/utils/tests/test_jsonutil.py index 45247a0..f70e941 100644 --- a/IPython/utils/tests/test_jsonutil.py +++ b/IPython/utils/tests/test_jsonutil.py @@ -12,12 +12,15 @@ #----------------------------------------------------------------------------- # stdlib import json +from base64 import decodestring # third party import nose.tools as nt # our own -from ..jsonutil import json_clean +from IPython.testing import decorators as dec +from ..jsonutil import json_clean, encode_images +from ..py3compat import unicode_to_str, str_to_bytes #----------------------------------------------------------------------------- # Test functions @@ -56,6 +59,35 @@ def test(): json.loads(json.dumps(out)) + +@dec.parametric +def test_encode_images(): + # invalid data, but the header and footer are from real files + pngdata = b'\x89PNG\r\n\x1a\nblahblahnotactuallyvalidIEND\xaeB`\x82' + jpegdata = b'\xff\xd8\xff\xe0\x00\x10JFIFblahblahjpeg(\xa0\x0f\xff\xd9' + + fmt = { + 'image/png' : pngdata, + 'image/jpeg' : jpegdata, + } + encoded = encode_images(fmt) + for key, value in fmt.iteritems(): + # encoded has unicode, want bytes + decoded = decodestring(encoded[key].encode('ascii')) + yield nt.assert_equal(decoded, value) + encoded2 = encode_images(encoded) + yield nt.assert_equal(encoded, encoded2) + + b64_str = {} + for key, encoded in encoded.iteritems(): + b64_str[key] = unicode_to_str(encoded) + encoded3 = encode_images(b64_str) + yield nt.assert_equal(encoded3, b64_str) + for key, value in fmt.iteritems(): + # encoded3 has str, want bytes + decoded = decodestring(str_to_bytes(encoded3[key])) + yield nt.assert_equal(decoded, value) + def test_lambda(): jc = json_clean(lambda : 1) assert isinstance(jc, str) diff --git a/IPython/zmq/displayhook.py b/IPython/zmq/displayhook.py index a3d9178..899b7ac 100644 --- a/IPython/zmq/displayhook.py +++ b/IPython/zmq/displayhook.py @@ -1,8 +1,8 @@ import __builtin__ import sys -from base64 import encodestring from IPython.core.displayhook import DisplayHook +from IPython.utils.jsonutil import encode_images from IPython.utils.traitlets import Instance, Dict from session import extract_header, Session @@ -30,18 +30,6 @@ class ZMQDisplayHook(object): self.parent_header = extract_header(parent) -def _encode_binary(format_dict): - encoded = format_dict.copy() - pngdata = format_dict.get('image/png') - if isinstance(pngdata, bytes): - encoded['image/png'] = encodestring(pngdata).decode('ascii') - jpegdata = format_dict.get('image/jpeg') - if isinstance(jpegdata, bytes): - encoded['image/jpeg'] = encodestring(jpegdata).decode('ascii') - - return encoded - - class ZMQShellDisplayHook(DisplayHook): """A displayhook subclass that publishes data using ZeroMQ. This is intended to work with an InteractiveShell instance. It sends a dict of different @@ -64,7 +52,7 @@ class ZMQShellDisplayHook(DisplayHook): self.msg['content']['execution_count'] = self.prompt_count def write_format_data(self, format_dict): - self.msg['content']['data'] = _encode_binary(format_dict) + self.msg['content']['data'] = encode_images(format_dict) def finish_displayhook(self): """Finish up all displayhook activities.""" diff --git a/IPython/zmq/zmqshell.py b/IPython/zmq/zmqshell.py index aedce0e..fe05bd2 100644 --- a/IPython/zmq/zmqshell.py +++ b/IPython/zmq/zmqshell.py @@ -38,12 +38,12 @@ from IPython.lib.kernel import ( ) from IPython.testing.skipdoctest import skip_doctest from IPython.utils import io -from IPython.utils.jsonutil import json_clean +from IPython.utils.jsonutil import json_clean, encode_images from IPython.utils.process import arg_split from IPython.utils import py3compat from IPython.utils.traitlets import Instance, Type, Dict, CBool, CBytes from IPython.utils.warn import warn, error -from IPython.zmq.displayhook import ZMQShellDisplayHook, _encode_binary +from IPython.zmq.displayhook import ZMQShellDisplayHook from IPython.zmq.session import extract_header from session import Session @@ -75,7 +75,7 @@ class ZMQDisplayPublisher(DisplayPublisher): self._validate_data(source, data, metadata) content = {} content['source'] = source - content['data'] = _encode_binary(data) + content['data'] = encode_images(data) content['metadata'] = metadata self.session.send( self.pub_socket, u'display_data', json_clean(content),