refactored nbconvert, nosetest pass

ivanov · ivanov · commit f1fa998c3980 · 2012-03-13T15:06:49.000-07:00
joint work with @jeremiahbuddha and @aisipos
diff --git a/nbconvert.py b/nbconvert.py
@@ -27,117 +27,141 @@ def unknown_cell(cell):
     return rst_directive('.. warning:: Unknown cell') + \
       [repr(cell)]
 
-def heading_cell(cell):
-    """convert a heading cell to rst
-
-    Returns list."""
-    heading_level = {1:'=', 2:'-', 3:'`', 4:'\'', 5:'.',6:'~'}
-    marker = heading_level[cell.level]
-    return ['{0}\n{1}\n'.format(cell.source, marker*len(cell.source))]
-
-def markdown_cell(cell):
-    """convert a markdown cell to rst
-
-    Returns list."""
-    return [cell.source]
-
 
 def rst_directive(directive, text=''):
     out = [directive, '']
     if text:
         out.extend([indent(text), ''])
     return out
 
-def code_cell(cell):
-    """Convert a code cell to rst
-
-    Returns list."""
-
-    if not cell.input:
-        return []
-
-    lines = ['In[%s]:' % cell.prompt_number, '']
-    lines.extend(rst_directive('.. code:: python', cell.input))
-    
-    for output in cell.outputs:
-        conv = converters.get(output.output_type, unknown_cell)
-        lines.extend(conv(output))
-
-    return lines
-
 # Converters for parts of a cell.
 figures_counter = 1
 
-def out_display(output):
-    """convert display data from the output of a code cell to rst.
+class ConversionException(Exception):
+    pass
+
+class Converter(object):
+    default_encoding = 'utf-8'
+    def __init__(self, fname):
+        self.fname = fname
+
+    @property
+    def extension(self):
+        raise ConversionException("""extension must be defined in Converter
+                subclass""")
+
+    def dispatch(self,cell_type):
+        """return cell_type dependent render method,  for example render_code
+        """
+        return getattr(self, 'render_'+cell_type, unknown_cell)
+
+    def convert(self):
+        lines = []
+        for cell in self.nb.worksheets[0].cells:
+            conv_fn = self.dispatch(cell.cell_type)
+            lines.extend(conv_fn(cell))
+            lines.append('')
+        return '\n'.join(lines)
+
+    def render(self):
+        "read, convert, and save self.fname"
+        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:
+            self.nb = nbformat.read(f, 'json')
+
+    def save(self,fname=None, encoding=None):
+        "read and parse notebook into self.nb"
+        if fname is None:
+            fname = os.path.splitext(self.fname)[0] + '.' + self.extension
+        if encoding is None:
+            encoding = self.default_encoding
+        with open(fname, 'w') as f:
+            f.write(self.output.encode(encoding))
+        return fname
 
-    Returns list.
-    """
-    global figures_counter
+    def render_heading(self,cell):
+         raise NotImplementedError
+    def render_code(self,cell):
+         raise NotImplementedError
+    def render_markdown(self,cell):
+         raise NotImplementedError
+    def render_pyout(self,cell):
+         raise NotImplementedError
+    def render_display_data(self,cell):
+         raise NotImplementedError
 
-    lines = []
+class ConverterRST(Converter):
+    extension = 'rst'
+    def render_heading(self,cell):
+        """convert a heading cell to rst
 
-    if 'png' in output:
-        fname = 'nb_figure_%s.png' % figures_counter
-        with open(fname, 'w') as f:
-            f.write(output.png.decode('base64'))
+        Returns list."""
+        heading_level = {1:'=', 2:'-', 3:'`', 4:'\'', 5:'.',6:'~'}
+        marker = heading_level[cell.level]
+        return ['{0}\n{1}\n'.format(cell.source, marker*len(cell.source))]
 
-        figures_counter += 1
-        lines.append('.. image:: %s' % fname)
-        lines.append('')
-    
-    return lines
+    def render_code(self,cell):
+        """Convert a code cell to rst
 
-    
-def out_pyout(output):
-    """convert pyout part of a code cell to rst
+        Returns list."""
 
-    Returns list."""
+        if not cell.input:
+            return []
 
-    lines = ['Out[%s]:' % output.prompt_number, '']
-    
-    if 'latex' in output:
-        lines.extend(rst_directive('.. math::', output.latex))
+        lines = ['In[%s]:' % cell.prompt_number, '']
+        lines.extend(rst_directive('.. code:: python', cell.input))
 
-    if 'text' in output:
-        lines.extend(rst_directive('.. parsed-literal::', output.text))
+        for output in cell.outputs:
+            conv_fn = self.dispatch(output.output_type)
+            lines.extend(conv_fn(output))
 
-    return lines
+        return lines
 
+    def render_markdown(self,cell):
+        """convert a markdown cell to rst
 
-converters = dict(heading = heading_cell,
-                  code = code_cell,
-                  markdown = markdown_cell,
-                  pyout = out_pyout,
-                  display_data = out_display,
-    )
+        Returns list."""
+        return [cell.source]
 
+    def render_pyout(self,output):
+        """convert pyout part of a code cell to rst
 
+        Returns list."""
 
-def convert_notebook(nb):
-    lines = []
-    for cell in nb.worksheets[0].cells:
-        conv = converters.get(cell.cell_type, unknown_cell)
-        lines.extend(conv(cell))
-        lines.append('')
-                
-    return '\n'.join(lines)
+        lines = ['Out[%s]:' % output.prompt_number, '']
+        
+        if 'latex' in output:
+            lines.extend(rst_directive('.. math::', output.latex))
 
+        if 'text' in output:
+            lines.extend(rst_directive('.. parsed-literal::', output.text))
 
-def nb2rst(fname):
-    "Convert notebook to rst"
-    
-    with open(fname) as f:
-        nb = nbformat.read(f, 'json')
+        return lines
 
-    rst = convert_notebook(nb)
+    def render_display_data(self,output):
+        """convert display data from the output of a code cell to rst.
 
-    newfname = os.path.splitext(fname)[0] + '.rst'
-    with open(newfname, 'w') as f:
-        f.write(rst.encode('utf8'))
+        Returns list.
+        """
+        global figures_counter
 
-    return newfname
+        lines = []
+
+        if 'png' in output:
+            fname = 'nb_figure_%s.png' % figures_counter
+            with open(fname, 'w') as f:
+                f.write(output.png.decode('base64'))
 
+            figures_counter += 1
+            lines.append('.. image:: %s' % fname)
+            lines.append('')
+        
+        return lines
 
 def rst2simplehtml(fname):
     """Convert a rst file to simplified html suitable for blogger.
@@ -191,7 +215,7 @@ def rst2simplehtml(fname):
 def main(fname):
     """Convert a notebook to html in one step"""
     newfname = nb2rst(fname)
-    rst2simplehtml(newfname)
+    #rst2simplehtml(newfname)
 
 
 if __name__ == '__main__':
diff --git a/tests/test.ipynb b/tests/test.ipynb
diff --git a/tests/test_simple.py b/tests/test_simple.py
@@ -0,0 +1,19 @@
+from nbconvert import ConverterRST
+import nose.tools as nt
+
+import os
+import glob
+
+def clean_dir():
+    "Remove .rst files created during conversion"
+    map(os.remove, glob.glob("*.rst"))
+    map(os.remove, glob.glob("*.png"))
+
+
+@nt.with_setup(clean_dir, clean_dir)
+def test_simple():
+    fname = 'test.ipynb'
+    c = ConverterRST(fname)
+    f = c.render()
+    nt.assert_true('rst' in f, 'changed file extension to rst')
+
