Extract output preprocessor only extracts specified formats

IPython/html/nbconvert/tests/test_nbconvert_handlers.py

@@ -1,114 +1,120 b''
 # coding: utf-8
+import base64
 import io
 import json
 import os
 from os.path import join as pjoin
 import shutil
 
 import requests
 
 from IPython.html.utils import url_path_join
 from IPython.html.tests.launchnotebook import NotebookTestBase, assert_http_error
 from IPython.nbformat.current import (new_notebook, write, new_worksheet,
                                       new_heading_cell, new_code_cell,
                                       new_output)
 
 class NbconvertAPI(object):
     """Wrapper for nbconvert API calls."""
     def __init__(self, base_url):
         self.base_url = base_url
 
     def _req(self, verb, path, body=None, params=None):
         response = requests.request(verb,
                 url_path_join(self.base_url, 'nbconvert', path),
                 data=body, params=params,
         )
         response.raise_for_status()
         return response
 
     def from_file(self, format, path, name, download=False):
         return self._req('GET', url_path_join(format, path, name),
                          params={'download':download})
 
     def from_post(self, format, nbmodel):
         body = json.dumps(nbmodel)
         return self._req('POST', format, body)
 
     def list_formats(self):
         return self._req('GET', '')
 
+png_green_pixel = base64.encodestring(b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00'
+b'\x00\x00\x01\x00\x00x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDAT'
+b'\x08\xd7c\x90\xfb\xcf\x00\x00\x02\\\x01\x1e.~d\x87\x00\x00\x00\x00IEND\xaeB`\x82')
+
 class APITest(NotebookTestBase):
     def setUp(self):
         nbdir = self.notebook_dir.name
         
         if not os.path.isdir(pjoin(nbdir, 'foo')):
             os.mkdir(pjoin(nbdir, 'foo'))
         
         nb = new_notebook(name='testnb')
         
         ws = new_worksheet()
         nb.worksheets = [ws]
         ws.cells.append(new_heading_cell(u'Created by test ³'))
         cc1 = new_code_cell(input=u'print(2*6)')
         cc1.outputs.append(new_output(output_text=u'12'))
+        cc1.outputs.append(new_output(output_png=png_green_pixel, output_type='pyout'))
         ws.cells.append(cc1)
         
         with io.open(pjoin(nbdir, 'foo', 'testnb.ipynb'), 'w',
                      encoding='utf-8') as f:
             write(nb, f, format='ipynb')
 
         self.nbconvert_api = NbconvertAPI(self.base_url())
 
     def tearDown(self):
         nbdir = self.notebook_dir.name
 
         for dname in ['foo']:
             shutil.rmtree(pjoin(nbdir, dname), ignore_errors=True)
     
     def test_from_file(self):
         r = self.nbconvert_api.from_file('html', 'foo', 'testnb.ipynb')
         self.assertEqual(r.status_code, 200)
         self.assertIn(u'text/html', r.headers['Content-Type'])
         self.assertIn(u'Created by test', r.text)
         self.assertIn(u'print', r.text)
 
         r = self.nbconvert_api.from_file('python', 'foo', 'testnb.ipynb')
         self.assertIn(u'text/x-python', r.headers['Content-Type'])
         self.assertIn(u'print(2*6)', r.text)
 
     def test_from_file_404(self):
         with assert_http_error(404):
             self.nbconvert_api.from_file('html', 'foo', 'thisdoesntexist.ipynb')
 
     def test_from_file_download(self):
         r = self.nbconvert_api.from_file('python', 'foo', 'testnb.ipynb', download=True)
         content_disposition = r.headers['Content-Disposition']
         self.assertIn('attachment', content_disposition)
         self.assertIn('testnb.py', content_disposition)
 
     def test_from_file_zip(self):
         r = self.nbconvert_api.from_file('latex', 'foo', 'testnb.ipynb', download=True)
         self.assertIn(u'application/zip', r.headers['Content-Type'])
         self.assertIn(u'.zip', r.headers['Content-Disposition'])
 
     def test_from_post(self):
         nbmodel_url = url_path_join(self.base_url(), 'api/notebooks/foo/testnb.ipynb')
         nbmodel = requests.get(nbmodel_url).json()
         
         r = self.nbconvert_api.from_post(format='html', nbmodel=nbmodel)
         self.assertEqual(r.status_code, 200)
         self.assertIn(u'text/html', r.headers['Content-Type'])
         self.assertIn(u'Created by test', r.text)
         self.assertIn(u'print', r.text)
         
         r = self.nbconvert_api.from_post(format='python', nbmodel=nbmodel)
         self.assertIn(u'text/x-python', r.headers['Content-Type'])
         self.assertIn(u'print(2*6)', r.text)
 
     def test_from_post_zip(self):
         nbmodel_url = url_path_join(self.base_url(), 'api/notebooks/foo/testnb.ipynb')
         nbmodel = requests.get(nbmodel_url).json()
 
         r = self.nbconvert_api.from_post(format='latex', nbmodel=nbmodel)
         self.assertIn(u'application/zip', r.headers['Content-Type'])
         self.assertIn(u'.zip', r.headers['Content-Disposition'])

IPython/nbconvert/preprocessors/extractoutput.py

@@ -1,102 +1,103 b''
 """Module containing a preprocessor that extracts all of the outputs from the
 notebook file.  The extracted outputs are returned in the 'resources' dictionary.
 """
 #-----------------------------------------------------------------------------
 # Copyright (c) 2013, the IPython Development Team.
 #
 # Distributed under the terms of the Modified BSD License.
 #
 # The full license is in the file COPYING.txt, distributed with this software.
 #-----------------------------------------------------------------------------
 
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 
 import base64
 import sys
 import os
 
-from IPython.utils.traitlets import Unicode
+from IPython.utils.traitlets import Unicode, Set
 from .base import Preprocessor
 from IPython.utils import py3compat
 
 #-----------------------------------------------------------------------------
 # Classes
 #-----------------------------------------------------------------------------
 
 class ExtractOutputPreprocessor(Preprocessor):
     """
     Extracts all of the outputs from the notebook file.  The extracted 
     outputs are returned in the 'resources' dictionary.
     """
 
     output_filename_template = Unicode(
         "{unique_key}_{cell_index}_{index}.{extension}", config=True)
 
+    extract_output_types = Set({'png', 'jpg', 'svg', 'pdf'}, config=True)
 
     def preprocess_cell(self, cell, resources, cell_index):
         """
         Apply a transformation on each cell,
         
         Parameters
         ----------
         cell : NotebookNode cell
             Notebook cell being processed
         resources : dictionary
             Additional resources used in the conversion process.  Allows
             preprocessors to pass variables into the Jinja engine.
         cell_index : int
             Index of the cell being processed (see base.py)
         """
 
         #Get the unique key from the resource dict if it exists.  If it does not 
         #exist, use 'output' as the default.  Also, get files directory if it
         #has been specified
         unique_key = resources.get('unique_key', 'output')
         output_files_dir = resources.get('output_files_dir', None)
         
         #Make sure outputs key exists
         if not isinstance(resources['outputs'], dict):
             resources['outputs'] = {}
             
         #Loop through all of the outputs in the cell
         for index, out in enumerate(cell.get('outputs', [])):
 
-            #Get the output in data formats that the template is interested in.
-            for out_type in self.display_data_priority:
+            #Get the output in data formats that the template needs extracted
+            for out_type in self.extract_output_types:
                 if out.hasattr(out_type): 
                     data = out[out_type]
 
                     #Binary files are base64-encoded, SVG is already XML
                     if out_type in ('png', 'jpg', 'jpeg', 'pdf'):
 
                         # data is b64-encoded as text (str, unicode)
                         # decodestring only accepts bytes
                         data = py3compat.cast_bytes(data)
                         data = base64.decodestring(data)
                     elif sys.platform == 'win32':
                         data = data.replace('\n', '\r\n').encode("UTF-8")
                     else:
                         data = data.encode("UTF-8")
                     
                     #Build an output name
                     filename = self.output_filename_template.format( 
                                     unique_key=unique_key,
                                     cell_index=cell_index,
                                     index=index,
                                     extension=out_type)
 
                     #On the cell, make the figure available via 
                     #   cell.outputs[i].svg_filename  ... etc (svg in example)
                     # Where
                     #   cell.outputs[i].svg  contains the data
                     if output_files_dir is not None:
                         filename = os.path.join(output_files_dir, filename)
                     out[out_type + '_filename'] = filename
 
                     #In the resources, make the figure available via
                     #   resources['outputs']['filename'] = data
                     resources['outputs'][filename] = data
 
         return cell, resources

IPython/nbconvert/templates/python.tpl

@@ -1,51 +1,51 b''
 {%- extends 'null.tpl' -%}
 
 
 {% block in_prompt %}
 # In[{{ cell.prompt_number if cell.prompt_number else ' ' }}]:
 {% endblock in_prompt %}
 
 {% block output_prompt %}
 # Out[{{ cell.prompt_number }}]:
 {% endblock output_prompt %}
 
 {% block input %}
 {{ cell.input | ipython2python }}
 {% endblock input %}
 
 {# Those Two are for error displaying
 even if the first one seem to do nothing, 
 it introduces a new line
 #}
 {% block pyerr %}
 {{ super() }}
 {% endblock pyerr %}
 
 {% block traceback_line %}
 {{ line | indent | strip_ansi }}
 {% endblock traceback_line %}
 {# .... #}
 
 {% block pyout %}
-{{ output.text | indent | comment_lines }}
+{{ output.text or '' | indent | comment_lines }}
 {% endblock pyout %}
 
 {% block stream %}
 {{ output.text | indent | comment_lines }}
 {% endblock stream %}
 
 {% block display_data scoped %}
 # image file:
 {% endblock display_data %}
 
 {% block markdowncell scoped %}
 {{ cell.source | comment_lines }}
 {% endblock markdowncell %}
 
 {% block headingcell scoped %}
 {{ '#' * cell.level }}{{ cell.source | replace('\n', ' ') | comment_lines }}
 {% endblock headingcell %}
 
 {% block unknowncell scoped %}
 unknown type  {{ cell.type }}
-{% endblock unknowncell %}
  No newline at end of file
+{% endblock unknowncell %}