Skip to content

Instantly share code, notes, and snippets.

@timo
Forked from minrk/checkipnb.py
Created May 6, 2012 11:10
Show Gist options
  • Save timo/2621679 to your computer and use it in GitHub Desktop.
Save timo/2621679 to your computer and use it in GitHub Desktop.
collect ipython notebook files for py.test runs
# put this file into any folder under which ipynb files shall be collected.
import pytest
import os,sys
wrapped_stdin = sys.stdin
sys.stdin = sys.__stdin__
from IPython.zmq.blockingkernelmanager import BlockingKernelManager
sys.stdin = wrapped_stdin
from IPython.nbformat.current import reads
# combined from
# http://pytest.org/latest/example/nonpython.html#non-python-tests
# and
# https://gist.github.com/2621679 by minrk
def pytest_collect_file(path, parent):
if path.ext == ".ipynb":
return IPyNbFile(path, parent)
class IPyNbFile(pytest.File):
def collect(self):
with self.fspath.open() as f:
self.nb = reads(f.read(), 'json')
cell_num = 0
for ws in self.nb.worksheets:
for cell in ws.cells:
if cell.cell_type == "code":
yield IPyNbCell(self.name, self, cell_num, cell)
cell_num += 1
def setup(self):
self.km = BlockingKernelManager()
self.km.start_kernel(stderr=open(os.devnull, 'w'))
self.km.start_channels()
self.shell = self.km.shell_channel
def teardown(self):
self.km.shutdown_kernel()
del self.shell
del self.km
class IPyNbCell(pytest.Item):
def __init__(self, name, parent, cell_num, cell):
super(IPyNbCell, self).__init__(name, parent)
self.cell_num = cell_num
self.cell = cell
def runtest(self):
shell = self.parent.shell
shell.execute(self.cell.input, allow_stdin=False)
# wait for finish, maximum 20s
reply = shell.get_msg(timeout=20)['content']
if reply['status'] == 'error':
raise IPyNbException(self.cell_num, self.cell.input, '\n'.join(reply['traceback']))
def repr_failure(self, excinfo):
""" called when self.runtest() raises an exception. """
if isinstance(excinfo.value, IPyNbException):
return "\n".join([
"notebook worksheet execution failed",
" cell %d\n\n"
" input: %r\n\n"
" raised: %r\n" % excinfo.value.args[1:3],
])
def reportinfo(self):
return self.fspath, 0, "cell %d" % self.cell_num
class IPyNbException(Exception):
""" custom exception for error reporting. """
@jakubLangr
Copy link

ImportError: No module named 'IPython.zmq'
Even though I did sudo pip install pyzmq tornado. IPyNB run on Unix.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment