Serve nbconvert output as zip when it has multiple files

IPython/html/nbconvert/handlers.py

@@ -1,4 +1,6 @@
+import io
 import os
+import zipfile
 
 from tornado import web
 
@@ -6,12 +8,44 @@ from ..base.handlers import IPythonHandler, notebook_path_regex
 from IPython.nbformat.current import to_notebook_json
 from IPython.nbconvert.exporters.export import exporter_map
 from IPython.utils import tz
-
-
-def has_resource_files(resources):
-    output_files_dir = resources.get('output_files_dir', "")
-    return bool(os.path.isdir(output_files_dir) and \
-                    os.listdir(output_files_dir))
+from IPython.utils.py3compat import cast_bytes
+
+import sys
+
+def find_resource_files(output_files_dir):
+    files = []
+    for dirpath, dirnames, filenames in os.walk(output_files_dir):
+        files.extend([os.path.join(dirpath, f) for f in filenames])
+    return files
+
+def respond_zip(handler, name, output, resources):
+    """Zip up the output and resource files and respond with the zip file.
+
+    Returns True if it has served a zip file, False if there are no resource
+    files, in which case we serve the plain output file.
+    """
+    # Check if we have resource files we need to zip
+    output_files = resources.get('outputs', None)
+    if not output_files:
+        return False
+
+    # Headers
+    zip_filename = os.path.splitext(name)[0] + '.zip'
+    handler.set_header('Content-Disposition',
+                       'attachment; filename="%s"' % zip_filename)
+    handler.set_header('Content-Type', 'application/zip')
+
+    # Prepare the zip file
+    buffer = io.BytesIO()
+    zipf = zipfile.ZipFile(buffer, mode='w', compression=zipfile.ZIP_DEFLATED)
+    output_filename = os.path.splitext(name)[0] + '.' + resources['output_extension']
+    zipf.writestr(output_filename, cast_bytes(output, 'utf-8'))
+    for filename, data in output_files.items():
+        zipf.writestr(os.path.basename(filename), data)
+    zipf.close()
+
+    handler.finish(buffer.getvalue())
+    return True
 
 class NbconvertFileHandler(IPythonHandler):
 
@@ -29,22 +63,22 @@ class NbconvertFileHandler(IPythonHandler):
         info = os.stat(os_path)
         self.set_header('Last-Modified', tz.utcfromtimestamp(info.st_mtime))
 
+        output, resources = exporter.from_filename(os_path)
+
+        if respond_zip(self, name, output, resources):
+            return
+
         # Force download if requested
         if self.get_argument('download', 'false').lower() == 'true':
-            filename = os.path.splitext(name)[0] + '.' + exporter.file_extension
+            filename = os.path.splitext(name)[0] + '.' + resources['output_extension']
             self.set_header('Content-Disposition',
-                            'attachment; filename="%s"' % filename)
-        
+                               'attachment; filename="%s"' % filename)
+
         # MIME type
         if exporter.output_mimetype:
             self.set_header('Content-Type',
                             '%s; charset=utf-8' % exporter.output_mimetype)
 
-        output, resources = exporter.from_filename(os_path)
-        
-        # TODO: If there are resources, combine them into a zip file
-        assert not has_resource_files(resources)
-        
         self.finish(output)
 
 class NbconvertPostHandler(IPythonHandler):
@@ -56,17 +90,17 @@ class NbconvertPostHandler(IPythonHandler):
         
         model = self.get_json_body()
         nbnode = to_notebook_json(model['content'])
-        
+
+        output, resources = exporter.from_notebook_node(nbnode)
+
+        if respond_zip(self, nbnode.metadata.name, output, resources):
+            return
+
         # MIME type
         if exporter.output_mimetype:
             self.set_header('Content-Type',
-            '%s; charset=utf-8' % exporter.output_mimetype)
-
-        output, resources = exporter.from_notebook_node(nbnode)
+                            '%s; charset=utf-8' % exporter.output_mimetype)
 
-        # TODO: If there are resources, combine them into a zip file
-        assert not has_resource_files(resources)
-        
         self.finish(output)
 
 #-----------------------------------------------------------------------------

IPython/html/nbconvert/tests/test_nbconvert_handlers.py

@@ -10,7 +10,8 @@ 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_heading_cell, new_code_cell,
+                                      new_output)
 
 class NbconvertAPI(object):
     """Wrapper for nbconvert API calls."""
@@ -48,7 +49,9 @@ class APITest(NotebookTestBase):
         ws = new_worksheet()
         nb.worksheets = [ws]
         ws.cells.append(new_heading_cell(u'Created by test ³'))
-        ws.cells.append(new_code_cell(input=u'print(2*6)'))
+        cc1 = new_code_cell(input=u'print(2*6)')
+        cc1.outputs.append(new_output(output_text=u'12'))
+        ws.cells.append(cc1)
         
         with io.open(pjoin(nbdir, 'foo', 'testnb.ipynb'), 'w',
                      encoding='utf-8') as f:
@@ -83,6 +86,11 @@ class APITest(NotebookTestBase):
         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()
@@ -96,3 +104,11 @@ class APITest(NotebookTestBase):
         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'])