Add argument parsing, and ability to convert an HTML file from command line

aisipos · aisipos · commit 0e8f2cf00171 · 2012-03-14T16:08:21.000-07:00
diff --git a/README.rst b/README.rst
@@ -24,6 +24,11 @@ The latest development version of doctest is required. This can be installed via
     $ curl http://docutils.svn.sourceforge.net/viewvc/docutils/trunk/docutils/?view=tar > docutils.gz
     $ pip install -U docutils.gz
 
+For conversion to HTML, pygments is also required
+::
+
+    $ pip install pygments
+
 Running Tests
 =============
 ::
diff --git a/nbconvert.py b/nbconvert.py
@@ -13,7 +13,7 @@
 import os
 import subprocess
 import sys
-
+import argparse
 from IPython.nbformat import current as nbformat
 from IPython.utils.text import indent
 
@@ -36,16 +36,17 @@ def rst_directive(directive, text=''):
 
 # Converters for parts of a cell.
 
+
 class ConversionException(Exception):
     pass
 
 
 class Converter(object):
     default_encoding = 'utf-8'
 
-    def __init__(self, fname):
-        self.fname = fname
-        self.dirpath = os.path.dirname(fname)
+    def __init__(self, infile):
+        self.infile = infile
+        self.dirpath = os.path.dirname(infile)
 
     @property
     def extension(self):
@@ -66,25 +67,25 @@ def convert(self):
         return '\n'.join(lines)
 
     def render(self):
-        "read, convert, and save self.fname"
+        "read, convert, and save self.infile"
         self.read()
         self.output = self.convert()
         return self.save()
 
     def read(self):
         "read and parse notebook into NotebookNode called self.nb"
-        with open(self.fname) as f:
+        with open(self.infile) as f:
             self.nb = nbformat.read(f, 'json')
 
-    def save(self, fname=None, encoding=None):
+    def save(self, infile=None, encoding=None):
         "read and parse notebook into self.nb"
-        if fname is None:
-            fname = os.path.splitext(self.fname)[0] + '.' + self.extension
+        if infile is None:
+            infile = os.path.splitext(self.infile)[0] + '.' + self.extension
         if encoding is None:
             encoding = self.default_encoding
-        with open(fname, 'w') as f:
+        with open(infile, 'w') as f:
             f.write(self.output.encode(encoding))
-        return fname
+        return infile
 
     def render_heading(self, cell):
         raise NotImplementedError
@@ -170,13 +171,13 @@ def render_display_data(self, output):
         lines = []
 
         if 'png' in output:
-            fname = 'nb_figure_%s.png' % self.figures_counter
-            fullname = os.path.join(self.dirpath, fname)
+            infile = 'nb_figure_%s.png' % self.figures_counter
+            fullname = os.path.join(self.dirpath, infile)
             with open(fullname, 'w') as f:
                 f.write(output.png.decode('base64'))
 
             self.figures_counter += 1
-            lines.append('.. image:: %s' % fname)
+            lines.append('.. image:: %s' % infile)
             lines.append('')
 
         return lines
@@ -194,7 +195,7 @@ def render_stream(self, output):
         return lines
 
 
-def rst2simplehtml(fname):
+def rst2simplehtml(infile):
     """Convert a rst file to simplified html suitable for blogger.
 
     This just runs rst2html with certain parameters to produce really simple
@@ -211,7 +212,7 @@ def rst2simplehtml(fname):
                     "--no-toc-backlinks --no-section-numbering "
                     "--strip-comments ")
 
-    cmd = "%s %s" % (cmd_template, fname)
+    cmd = "%s %s" % (cmd_template, infile)
     proc = subprocess.Popen(cmd,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE,
@@ -232,7 +233,7 @@ def rst2simplehtml(fname):
         if line.startswith('<body>'):
             break
 
-    newfname = os.path.splitext(fname)[0] + '.html'
+    newfname = os.path.splitext(infile)[0] + '.html'
     with open(newfname, 'w') as f:
         for line in walker:
             if line.startswith('</body>'):
@@ -243,11 +244,27 @@ def rst2simplehtml(fname):
     return newfname
 
 
-def main(fname):
+def main(infile, format='rst'):
     """Convert a notebook to html in one step"""
-    converter = ConverterRST(fname)
-    converter.render()
+    if format == 'rst':
+        converter = ConverterRST(infile)
+        converter.render()
+    elif format == 'html':
+        #Currently, conversion to html is a 2 step process, nb->rst->html
+        converter = ConverterRST(infile)
+        rstfname = converter.render()
+        rst2simplehtml(rstfname)
 
 
 if __name__ == '__main__':
-    main(sys.argv[1])
+    parser = argparse.ArgumentParser(description='nbconvert: Convert IPython notebooks to other formats')
+
+    # 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')
+    args = parser.parse_args()
+    main(infile=args.infile[0], format=args.format)
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -13,13 +13,16 @@ 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"))
+
 
 @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.with_setup(clean_dir, clean_dir)
 def test_main():
     """
@@ -28,14 +31,15 @@ def test_main():
     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):
+    for level in xrange(1, 7):
         cell = {
             'cell_type': 'heading',
             'level'    : level,
-            'source'   :  ['Test for heading type H{0}'.format(level)] 
+            'source'   :  ['Test for heading type H{0}'.format(level)]
             }
         # Convert cell dictionaries to NotebookNode
         cell_nb = nbformat.NotebookNode(cell)
@@ -48,10 +52,19 @@ def test_render_heading():
         # Render to rst
         c = ConverterRST('')
         rst_list = c.render_heading(cell_nb)
-        nt.assert_true(isinstance(rst_list,list)) # render should return a list
+        nt.assert_true(isinstance(rst_list, list))  # render should return a list
         rst_str = "".join(rst_list)
         # Confirm rst content
-        heading_level = {1:'=', 2:'-', 3:'`', 4:'\'', 5:'.',6:'~'}
+        heading_level = {1: '=', 2: '-', 3: '`', 4: '\'', 5: '.', 6: '~'}
         chk_str = "Test for heading type H{0}\n{1}\n".format(
-             level,heading_level[level]*24)
-        nt.assert_equal(rst_str,chk_str)
+             level, 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'))
