Merge pull request #57 from maxalbert/cleanup

Cleanup of format selection

Two small cleanups/fixes:

Much cleaner way of finding the right converter. This will result in much cleaner diffs when changing the API, adding arguments, etc. It also allows us to extract the known formats automatically.

Fixed an assertion which didn't actually check the file extension but an arbitrary occurrence of 'rst'

nbconvert.py

@@ -1,74 +1,78 b''
 #!/usr/bin/env python
 """Convert IPython notebooks to other formats, such as ReST, and HTML.
 
 Example:
   ./nbconvert.py --format rst file.ipynb
 
 Produces 'file.rst', along with auto-generated figure files
 called nb_figure_NN.png.
 """
 #-----------------------------------------------------------------------------
 # Imports
 #-----------------------------------------------------------------------------
 from __future__ import print_function
 
-
 # From IPython
 from IPython.external import argparse
 
 # local
 from converters.html import ConverterHTML
 from converters.markdown import ConverterMarkdown
 from converters.bloggerhtml import ConverterBloggerHTML
 from converters.rst import ConverterRST
 from converters.latex import ConverterLaTeX
 from converters.notebook import ConverterNotebook
 from converters.python import ConverterPy
 
-known_formats = "rst (default), html, blogger-html, latex, markdown, py"
+
+# When adding a new format, make sure to add it to the `converters`
+# dictionary below. This is used to create the list of known formats,
+# which gets printed in case an unknown format is encounteres, as well
+# as in the help
+
+converters = {
+    'rst': ConverterRST,
+    'markdown': ConverterMarkdown,
+    'html': ConverterHTML,
+    'blogger-html': ConverterBloggerHTML,
+    'latex': ConverterLaTeX,
+    'py': ConverterPy,
+    }
+
+default_format = 'rst'
+
+# Extract the list of known formats and mark the first format as the default.
+known_formats = ', '.join([key + " (default)" if key == default_format else key
+                           for key in converters])
+
 
 def main(infile, format='rst'):
     """Convert a notebook to html in one step"""
-    # XXX: this is just quick and dirty for now. When adding a new format,
-    # make sure to add it to the `known_formats` string above, which gets
-    # printed in in the catch-all else, as well as in the help
-    if format == 'rst':
-        converter = ConverterRST(infile)
-        converter.render()
-    elif format == 'markdown':
-        converter = ConverterMarkdown(infile)
-        converter.render()
-    elif format == 'html':
-        converter = ConverterHTML(infile)
-        converter.render()
-    elif format == 'blogger-html':
-        converter = ConverterBloggerHTML(infile)
-        converter.render()
-    elif format == 'latex':
-        converter = ConverterLaTeX(infile)
-        converter.render()
-    elif format == 'py':
-        converter = ConverterPy(infile)
-        converter.render()
-    else:
+
+    try:
+        ConverterClass = converters[format]
+    except KeyError:
         raise SystemExit("Unknown format '%s', " % format +
                 "known formats are: " + known_formats)
 
+    converter = ConverterClass(infile)
+    converter.render()
+
 #-----------------------------------------------------------------------------
 # Script main
 #-----------------------------------------------------------------------------
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description=__doc__,
             formatter_class=argparse.RawTextHelpFormatter)
     # TODO: consider passing file like object around, rather than filenames
     # would allow us to process stdin, or even http streams
     #parser.add_argument('infile', nargs='?', type=argparse.FileType('r'), default=sys.stdin)
 
     #Require a filename as a positional argument
     parser.add_argument('infile', nargs=1)
     parser.add_argument('-f', '--format', default='rst',
                         help='Output format. Supported formats: \n' +
                         known_formats)
     args = parser.parse_args()
     main(infile=args.infile[0], format=args.format)

tests/test_simple.py

@@ -1,71 +1,71 b''
 from nbconvert import ConverterRST, main
 import nose.tools as nt
 
 import os
 import glob
 from IPython.nbformat import current as nbformat
 
 fname = 'tests/test.ipynb'
 out_fname = 'tests/test.rst'
 
 
 def clean_dir():
     "Remove .rst files created during conversion"
     map(os.remove, glob.glob("./tests/*.rst"))
     map(os.remove, glob.glob("./tests/*.png"))
     map(os.remove, glob.glob("./tests/*.html"))
     map(os.remove, glob.glob("./tests/test_files/*"))
 
 
 @nt.with_setup(clean_dir, clean_dir)
 def test_simple():
     c = ConverterRST(fname)
     f = c.render()
-    nt.assert_true('rst' in f, 'changed file extension to rst')
+    nt.assert_true(f.endswith('.rst'), 'changed file extension to rst')
 
 
 @nt.with_setup(clean_dir, clean_dir)
 def test_main():
     """
     Test main entry point
     """
     main(fname)
     nt.assert_true(os.path.exists(out_fname))
 
 
 def test_render_heading():
     """ Unit test for cell type "heading" """
     # Generate and test heading cells level 1-6
     for level in xrange(1, 7):
         cell = {
             'cell_type': 'heading',
             'level': level,
             'source':  ['Test for heading type H{0}'.format(level)]
             }
         # Convert cell dictionaries to NotebookNode
         cell_nb = nbformat.NotebookNode(cell)
         # Make sure "source" attribute is uniconde not list.
         # For some reason, creating a NotebookNode manually like
         # this isn't converting source to a string like using
         # the create-from-file routine.
         if type(cell_nb.source) is list:
             cell_nb.source = '\n'.join(cell_nb.source)
         # Render to rst
         c = ConverterRST('')
         rst_list = c.render_heading(cell_nb)
         # render should return a list
         nt.assert_true(isinstance(rst_list, list))
         rst_str = "".join(rst_list)
         # Confirm rst content
         chk_str = "Test for heading type H{0}\n{1}\n".format(
              level, c.heading_level[level] * 24)
         nt.assert_equal(rst_str, chk_str)
 
 
 @nt.with_setup(clean_dir, clean_dir)
 def test_main_html():
     """
     Test main entry point
     """
     main(fname, format='html')
     nt.assert_true(os.path.exists('tests/test.html'))