Created
March 31, 2014 17:58
-
-
Save patricksnape/9898229 to your computer and use it in GitHub Desktop.
Nose plugin to run IPython notebooks as tests (run each notebook and check if any cells throw exceptions)
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import unittest | |
import nose | |
from IPython.kernel import KernelManager | |
from IPython.nbformat.current import reads | |
from nose.plugins.base import Plugin | |
from nose.util import anyp | |
class KernelProxy(object): | |
def __init__(self): | |
self.kernel_manager = None | |
self.client = None | |
self.shell = None | |
def __del__(self): | |
self.destroy_kernel() | |
def setup_kernel(self): | |
self.kernel_manager = KernelManager() | |
self.kernel_manager.start_kernel(stderr=open(os.devnull, 'w')) | |
try: | |
self.client = self.kernel_manager.client() | |
except AttributeError: | |
# 0.13 | |
self.client = self.kernel_manager | |
self.client.start_channels() | |
self.shell = self.client.shell_channel | |
def destroy_kernel(self): | |
try: | |
self.client.stop_channels() | |
self.kernel_manager.shutdown_kernel() | |
del self.kernel_manager | |
del self.client | |
del self.shell | |
except: # Eat all exceptions when destroying kernel | |
pass | |
class IPynbFileCase(unittest.TestCase): | |
_multiprocess_can_split_ = False | |
_multiprocess_shared_ = False | |
def __init__(self, cell_number, cell, notebook_name, kernel_proxy): | |
unittest.TestCase.__init__(self) | |
self.cell = cell | |
self.kernel_proxy = kernel_proxy | |
self.notebook_name = notebook_name | |
self.cell_number = cell_number | |
def id(self): | |
return '{0}[{1}]'.format(self.notebook_name, | |
self.cell_number) | |
def shortDescription(self): | |
return self.id() | |
def runTest(self): | |
self.kernel_proxy.shell.execute(self.cell.input) | |
# wait for finish, maximum 20s | |
reply = self.kernel_proxy.shell.get_msg(timeout=20)['content'] | |
if reply['status'] == 'error': | |
raise Exception('\n'.join(reply['traceback'])) | |
class IPythonNotebookTester(Plugin): | |
""" | |
Activate IPython Notebook tester | |
""" | |
extension = 'ipynb' | |
name = 'ipynb' | |
def __init__(self): | |
Plugin.__init__(self) | |
self.current_kernel = None | |
def options(self, parser, env): | |
"""Register commmandline options. | |
""" | |
Plugin.options(self, parser, env) | |
def configure(self, options, config): | |
"""Configure plugin. | |
""" | |
Plugin.configure(self, options, config) | |
def beforeContext(self): | |
del self.current_kernel | |
self.current_kernel = KernelProxy() | |
self.current_kernel.setup_kernel() | |
def afterContext(self): | |
if self.current_kernel is not None: | |
self.current_kernel.destroy_kernel() | |
self.current_kernel = None | |
def loadTestsFromFile(self, filename): | |
"""Load ipynb cells and run them as 'tests' | |
Tests are loaded only if filename's extension matches 'ipynb'. | |
""" | |
if self.is_notebook(filename): | |
basename = os.path.basename(filename) | |
with open(filename) as f: | |
nb = reads(f.read(), 'json') | |
suite = [] | |
cell_count = 0 | |
for ws in nb.worksheets: | |
for cell in ws.cells: | |
if cell.cell_type != 'code': | |
continue | |
suite.append(IPynbFileCase(cell_count, cell, basename, | |
self.current_kernel)) | |
cell_count += 1 | |
return suite | |
def is_notebook(self, name): | |
return anyp(name.endswith, self.extension) | |
def wantFile(self, file): | |
"""Override to select all modules and any file ending with 'ipynb'. | |
""" | |
if self.is_notebook(file): | |
return True | |
return None | |
if __name__ == '__main__': | |
nose.run('menpo', argv=['--with-ipynb', '-v'], | |
addplugins=[IPythonNotebookTester()]) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment