"""
Helper utility to PEP8 check Python code in IPython notebooks.
"""
import json
import os
import subprocess
import sys
import tempfile


TEMPLATE_IPYTHON_PRELUDE = """import numpy as np
"""


def find_notebooks(start_directory, extension="ipynb"):
    """
    Finds all notebooks in the start directory.
    """
    notebooks = []
    for root, _, files in os.walk(start_directory):
        notebooks = notebooks + [os.path.join(root, notebook)
                                 for notebook in files
                                 if notebook.endswith(extension)]
    return notebooks


def read_code_from_notebook(notebook):
    """
    Read the source code in a IPython notebook.
    """
    with open(notebook, 'r') as handle:
        text = handle.read()
        data = json.loads(text)
        cells = (cell['input']
                 for cell in data['worksheets'][0]['cells']
                 if cell['cell_type'] == u'code')

        return cells


def check_code(cells):
    """
    Check the code in the cells
    """
    temp = tempfile.NamedTemporaryFile()
    try:
        # Gather all code in a temp file together with a pylab prelude
        temp.write(TEMPLATE_IPYTHON_PRELUDE)
        for cell in cells:
            temp.write("".join(cell))

        process = subprocess.Popen(['pep8', temp.name],
                                   stdout=subprocess.PIPE)
        result, _ = process.communicate()
        print result
        if result == "":
            success = True
        else:
            success = False
    finally:
        temp.close()

    return success


def check_files(start_directory):
    """
    Checks all files under start directory.
    """
    notebooks = find_notebooks(start_directory)
    for notebook in notebooks:
        cells = read_code_from_notebook(notebook)
        success = check_code(cells)
        if not success:
            break

if __name__ == "__main__":
    if len(sys.argv) == 1:
        ROOT = "."
    else:
        ROOT = sys.argv[1]
    check_files(ROOT)

@@ -131,16 +131,50 @@
 empfile%0A
+from optparse import OptionParser%0A
 %0A%0ATEMPLA
@@ -1008,16 +1008,25 @@
 de(cells
+, options
 ):%0A    %22
@@ -1105,24 +1105,25 @@
 oraryFile()%0A
+%0A
     try:%0A   
@@ -1119,16 +1119,142 @@
    try:%0A
+        if options.filename is not None:%0A            handle = open(options.filename, 'w')%0A%0A        code = %22%22%0A        try:%0A    
         
@@ -1324,27 +1324,27 @@
         
-temp.write(
+    code = 
 TEMPLATE
@@ -1359,18 +1359,21 @@
 _PRELUDE
-)
 %0A
+    
         
@@ -1407,35 +1407,195 @@
     
-temp.write(%22%22.join(cell))%0A%0A
+    code += %22%22.join(cell) + %22%5Cn%5Cn%5Cn%22%0A            temp.write(code)%0A            temp.flush()%0A            if options.filename is not None and handle:%0A                handle.write(code)%0A%0A    
     
@@ -1681,16 +1681,20 @@
         
+    
 stdout=s
@@ -1709,16 +1709,20 @@
 s.PIPE)%0A
+    
         
@@ -1763,16 +1763,20 @@
         
+    
 print re
@@ -1788,16 +1788,20 @@
         
+   
+ 
 if resul
@@ -1813,32 +1813,36 @@
 %22%22:%0A            
+    
 success = True%0A 
@@ -1844,30 +1844,38 @@
 rue%0A        
+  
+  
 else:%0A
+    
             
@@ -1894,16 +1894,20 @@
 lse%0A    
+  
+  
 finally:
@@ -1907,16 +1907,20 @@
 inally:%0A
+    
         
@@ -1931,16 +1931,108 @@
 .close()
+%0A    finally:%0A        if options.filename is not None and handle:%0A            handle.close()
 %0A%0A    re
@@ -2073,24 +2073,33 @@
 rt_directory
+, options
 ):%0A    %22%22%22%0A 
@@ -2141,32 +2141,84 @@
 ectory.%0A    %22%22%22%0A
+    failfast = options.failfast%0A%0A    success = True%0A
     notebooks = 
@@ -2364,16 +2364,25 @@
 de(cells
+, options
 )%0A      
@@ -2386,16 +2386,29 @@
      if 
+failfast and 
 not succ
@@ -2429,16 +2429,35 @@
    break
+%0A    return success
 %0A%0Aif __n
@@ -2485,105 +2485,704 @@
     
-if len(sys.argv) == 1:%0A        ROOT = %22.%22%0A    else:%0A        ROOT = sys.argv%5B1%5D%0A    check_files(ROOT
+PARSER = OptionParser()%0A    PARSER.add_option(%22-o%22, %22--output%22, dest=%22filename%22,%0A                      help=%22write combined python code to FILE%22,%0A                      metavar=%22FILE%22)%0A    PARSER.add_option(%22-f%22, %22--fail-fast%22, dest=%22failfast%22, default=False,%0A                      action=%22store_true%22,%0A                      help=%22fail when the first notebook with an error%5C%0A                              is encountered%22)%0A    (OPTION, ARGS) = PARSER.parse_args()%0A%0A    if len(ARGS) == 0:%0A        SUCCESSES = list(check_files(%22.%22, OPTION))%0A    else:%0A        SUCCESSES = %5Bcheck_files(arg, OPTION) for arg in ARGS%5D%0A    if reduce(lambda x, y: x and y, SUCCESSES):%0A        exit(0)%0A    else:%0A        exit(1
 )%0A

Write complete file, add some options