move _encode_binary to jsonutil.encode_images

minrk · minrk · commit 001424863771 · 2012-06-27T19:24:08.000-07:00
And fix identification of PNG/JPEG data, which could be double-encoded when using a JSON library that decodes to str instead of unicode (simplejson 2.6, for instance).

Changes are tested.
diff --git 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
@@ -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
@@ -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 @@ def set_parent(self, parent):
         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 @@ def write_output_prompt(self):
         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
@@ -38,12 +38,12 @@
 )
 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 @@ def publish(self, source, data, metadata=None):
         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),
