Serve nbconvert output as zip when it has multiple files

This commit is contained in:
Thomas Kluyver 2013-12-19 16:07:27 -08:00
parent 57642b3847
commit be6a5d32ea
2 changed files with 72 additions and 22 deletions

View File

@ -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
from IPython.utils.py3compat import cast_bytes
import sys
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))
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'])
# 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)
# TODO: If there are resources, combine them into a zip file
assert not has_resource_files(resources)
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)
self.finish(output)
#-----------------------------------------------------------------------------

View File

@ -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'])