diff --git a/IPython/nbformat/v2/nbjson.py b/IPython/nbformat/v2/nbjson.py index 249439f..b86297d 100644 --- a/IPython/nbformat/v2/nbjson.py +++ b/IPython/nbformat/v2/nbjson.py @@ -16,9 +16,8 @@ Authors: # Imports #----------------------------------------------------------------------------- -from base64 import encodestring from .nbbase import from_dict -from .rwbase import NotebookReader, NotebookWriter, base64_decode +from .rwbase import NotebookReader, NotebookWriter, restore_bytes import json #----------------------------------------------------------------------------- @@ -26,9 +25,10 @@ import json #----------------------------------------------------------------------------- class BytesEncoder(json.JSONEncoder): + """A JSON encoder that accepts b64 (and other *ascii*) bytestrings.""" def default(self, obj): if isinstance(obj, bytes): - return encodestring(obj).decode('ascii') + return obj.decode('ascii') return json.JSONEncoder.default(self, obj) @@ -40,7 +40,7 @@ class JSONReader(NotebookReader): return nb def to_notebook(self, d, **kwargs): - return base64_decode(from_dict(d)) + return restore_bytes(from_dict(d)) class JSONWriter(NotebookWriter): diff --git a/IPython/nbformat/v2/rwbase.py b/IPython/nbformat/v2/rwbase.py index 2641639..cbcb50a 100644 --- a/IPython/nbformat/v2/rwbase.py +++ b/IPython/nbformat/v2/rwbase.py @@ -19,31 +19,67 @@ Authors: from base64 import encodestring, decodestring import pprint +from IPython.utils.py3compat import str_to_bytes + #----------------------------------------------------------------------------- # Code #----------------------------------------------------------------------------- +def restore_bytes(nb): + """Restore bytes of image data from unicode-only formats. + + Base64 encoding is handled elsewhere. Bytes objects in the notebook are + always b64-encoded. We DO NOT encode/decode around file formats. + """ + for ws in nb.worksheets: + for cell in ws.cells: + if cell.cell_type == 'code': + for output in cell.outputs: + if 'png' in output: + output.png = str_to_bytes(output.png, 'ascii') + if 'jpeg' in output: + output.jpeg = str_to_bytes(output.jpeg, 'ascii') + return nb + + +# b64 encode/decode are never actually used, because all bytes objects in +# the notebook are already b64-encoded, and we don't need/want to double-encode + def base64_decode(nb): - """Base64 encode all bytes objects in the notebook.""" + """Restore all bytes objects in the notebook from base64-encoded strings. + + Note: This is never used + """ for ws in nb.worksheets: for cell in ws.cells: if cell.cell_type == 'code': - if 'png' in cell: - cell.png = bytes(decodestring(cell.png)) - if 'jpeg' in cell: - cell.jpeg = bytes(decodestring(cell.jpeg)) + for output in cell.outputs: + if 'png' in output: + if isinstance(output.png, unicode): + output.png = output.png.encode('ascii') + output.png = decodestring(output.png) + if 'jpeg' in output: + if isinstance(output.jpeg, unicode): + output.jpeg = output.jpeg.encode('ascii') + output.jpeg = decodestring(output.jpeg) return nb def base64_encode(nb): - """Base64 decode all binary objects in the notebook.""" + """Base64 encode all bytes objects in the notebook. + + These will be b64-encoded unicode strings + + Note: This is never used + """ for ws in nb.worksheets: for cell in ws.cells: if cell.cell_type == 'code': - if 'png' in cell: - cell.png = unicode(encodestring(cell.png)) - if 'jpeg' in cell: - cell.jpeg = unicode(encodestring(cell.jpeg)) + for output in cell.outputs: + if 'png' in output: + output.png = encodestring(output.png).decode('ascii') + if 'jpeg' in output: + output.jpeg = encodestring(output.jpeg).decode('ascii') return nb diff --git a/IPython/nbformat/v2/tests/nbexamples.py b/IPython/nbformat/v2/tests/nbexamples.py index b567884..2f6174b 100644 --- a/IPython/nbformat/v2/tests/nbexamples.py +++ b/IPython/nbformat/v2/tests/nbexamples.py @@ -1,10 +1,15 @@ +import os +from base64 import encodestring + from ..nbbase import ( NotebookNode, new_code_cell, new_text_cell, new_worksheet, new_notebook, new_output, new_metadata, new_author ) - +# some random base64-encoded *bytes* +png = encodestring(os.urandom(5)) +jpeg = encodestring(os.urandom(6)) ws = new_worksheet(name='worksheet1') @@ -42,8 +47,8 @@ ws.cells.append(new_code_cell( output_text=u'', output_html=u'The HTML rep', output_latex=u'$a$', - output_png=b'data', - output_jpeg=b'data', + output_png=png, + output_jpeg=jpeg, output_svg=u'', output_json=u'json data', output_javascript=u'var i=0;', @@ -53,8 +58,8 @@ ws.cells.append(new_code_cell( output_text=u'', output_html=u'The HTML rep', output_latex=u'$a$', - output_png=b'data', - output_jpeg=b'data', + output_png=png, + output_jpeg=jpeg, output_svg=u'', output_json=u'json data', output_javascript=u'var i=0;'